From 188bf9bd1a0aff02550523afbe2e61f50d31c175 Mon Sep 17 00:00:00 2001 From: Logiquo Date: Wed, 28 Oct 2026 05:49:02 -0400 Subject: [PATCH 01/11] Add CrossMap.has_many_targets to tell x-to-one mappings from x-to-many Lets callers check up front whether map() can return more than one target code, e.g. to reject a code mapping that would duplicate records. The icd-mappings backend always returns a single target; PyHealth tables are checked over their entries. --- pyhealth/medcode/cross_map.py | 13 +++++++++++++ tests/core/test_medcode_icd_mappings.py | 17 +++++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/pyhealth/medcode/cross_map.py b/pyhealth/medcode/cross_map.py index 72f5f44f3..ee1773ff1 100644 --- a/pyhealth/medcode/cross_map.py +++ b/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})" diff --git a/tests/core/test_medcode_icd_mappings.py b/tests/core/test_medcode_icd_mappings.py index 1b24f21cd..ba5561b19 100644 --- a/tests/core/test_medcode_icd_mappings.py +++ b/tests/core/test_medcode_icd_mappings.py @@ -160,6 +160,23 @@ def test_native_backend_exposes_empty_unmapped_codes(self): self.assertEqual(mapping.unmapped_codes, set()) +class TestHasManyTargets(unittest.TestCase): + """`has_many_targets` tells x-to-one tables from x-to-many ones.""" + + def test_icd_mappings_backend_has_single_targets(self): + for pair in ICD_MAPPINGS_PAIRS: + with self.subTest(pair=pair): + self.assertFalse(CrossMap.load(*pair).has_many_targets) + + def test_pyhealth_table_with_several_targets_has_many_targets(self): + # Built by hand so PyHealth's resource server is never contacted. + mapping = CrossMap.__new__(CrossMap) + mapping.mapping = {"A": ["1"], "B": ["1"]} + self.assertFalse(mapping.has_many_targets) + mapping.mapping = {"A": ["1"], "B": ["1", "2"]} + self.assertTrue(mapping.has_many_targets) + + class TestBackendSelectionEndToEnd(unittest.TestCase): """The `backend` argument must survive the trip through load().""" From 1417039c9d4205e5ae09c155355863bb68e15b8a Mon Sep 17 00:00:00 2001 From: Logiquo Date: Wed, 28 Oct 2026 06:55:30 -0400 Subject: [PATCH 02/11] Add SequenceProcessorV2 for hierarchical packed records A single processor for nested codes, values and timestamps at any depth, emitted in CSR layout (payloads plus offsets) without padding or truncation. - Records: plain-dict columnar input (time, tokens, values, children). - LevelConfig: per-level ordering, token namespaces K, value shape S, code mapping, resampling and spatial axes; validated at construction. fit() only builds vocabularies. - Code mapping through CrossMap; multi-target mappings expand records and are only allowed on set levels without values. - Resampling of series levels onto a time grid with first, last, min, max, avg or merge; steps without records are left for imputation. - collate() and validate() for packed batches. - normalize, impute and missing_indicator are reserved for future work. --- docs/api/processors.rst | 1 + ...yhealth.processors.SequenceProcessorV2.rst | 54 ++ pyhealth/processors/__init__.py | 10 + pyhealth/processors/sequence_processor_v2.py | 882 ++++++++++++++++++ tests/core/test_sequence_processor_v2.py | 714 ++++++++++++++ 5 files changed, 1661 insertions(+) create mode 100644 docs/api/processors/pyhealth.processors.SequenceProcessorV2.rst create mode 100644 pyhealth/processors/sequence_processor_v2.py create mode 100644 tests/core/test_sequence_processor_v2.py diff --git a/docs/api/processors.rst b/docs/api/processors.rst index f4f2996fb..40010b826 100644 --- a/docs/api/processors.rst +++ b/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 diff --git a/docs/api/processors/pyhealth.processors.SequenceProcessorV2.rst b/docs/api/processors/pyhealth.processors.SequenceProcessorV2.rst new file mode 100644 index 000000000..0389b5588 --- /dev/null +++ b/docs/api/processors/pyhealth.processors.SequenceProcessorV2.rst @@ -0,0 +1,54 @@ +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``, ``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]``. ``normalize``, +``impute`` and ``missing_indicator`` are reserved for future work and only +accept ``None`` / ``False``; values 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: diff --git a/pyhealth/processors/__init__.py b/pyhealth/processors/__init__.py index b0738b05a..dcf368bee 100644 --- a/pyhealth/processors/__init__.py +++ b/pyhealth/processors/__init__.py @@ -37,6 +37,12 @@ def get_processor(name: str): ) from .raw_processor import RawProcessor from .sequence_processor import SequenceProcessor +from .sequence_processor_v2 import ( + LevelConfig, + Records, + Resample, + SequenceProcessorV2, +) from .signal_processor import SignalProcessor from .stagenet_processor import ( StageNetProcessor, @@ -76,6 +82,10 @@ def get_processor(name: str): "NestedSequenceProcessor", "RawProcessor", "SequenceProcessor", + "SequenceProcessorV2", + "Records", + "LevelConfig", + "Resample", "SignalProcessor", "StageNetProcessor", "StageNetTensorProcessor", diff --git a/pyhealth/processors/sequence_processor_v2.py b/pyhealth/processors/sequence_processor_v2.py new file mode 100644 index 000000000..96c64b2d7 --- /dev/null +++ b/pyhealth/processors/sequence_processor_v2.py @@ -0,0 +1,882 @@ +"""Unified processor for hierarchical packed records (CSR / jagged layout). + +See ``notes/Unified_Sequence_Processor.md`` (Section 7) for the design. +""" + +from collections.abc import Iterable, Sequence +from dataclasses import dataclass +from itertools import accumulate, repeat +from typing import TYPE_CHECKING, Any, Literal, TypedDict, get_args + +import numpy as np +import numpy.typing as npt +import torch + +from . import register_processor +from .base_processor import FeatureProcessor, TokenProcessorInterface + +if TYPE_CHECKING: + from pyhealth.medcode import CrossMap + +OrderType = Literal["set", "sequence", "series"] +Role = Literal["tokens", "values", "time", "offsets"] +FieldKey = tuple[int, Role] +CodeMapping = tuple[str, str] +ResampleMethod = Literal["first", "last", "min", "max", "avg", "merge"] + + +# Options reserved in the configuration but not implemented in Phase 1. +_RESERVED = {"normalize": None, "impute": None, "missing_indicator": False} + + +class Records(TypedDict, total=False): + """Raw input of one level: a collection of ``N`` records, stored as columns. + + A ``Records`` is a plain ``dict``; this class only provides type hints. + Level 0 is the innermost, and the one a task returns is the outermost + level ``L``. The level's :class:`LevelConfig` decides which keys are read; + a key it needs must be present (its column may be empty), and any other + key is ignored. + + Keys: + time: ``[N]`` numbers in the task's own unit. Read on ``series`` + levels. + tokens: ``[N]`` tokens, or ``[N]`` tuples of ``K`` tokens when the + level has ``K`` token namespaces. + values: ``[N, *S]`` numbers. Scalars are ``[N, 1]``; a flat ``[N]`` + list is also accepted when ``S`` is ``(1,)``. + children: ``[N]`` ``Records``; ``children[i]`` holds record ``i``'s + records one level down. Read on every level but 0. + """ + + time: Sequence[float] | np.ndarray + tokens: Sequence[str] | Sequence[tuple[str, ...]] + values: npt.ArrayLike + children: Sequence["Records"] + + +class Resample(TypedDict): + """Resampling of a ``series`` level onto a time grid; a plain ``dict``. + + Each grid step becomes one record, and the level becomes a ``sequence`` + without time: record ``i`` is the step starting at ``i * step``. The grid + starts at the task's time 0, so negative time raises. + + Keys: + step: Grid step, in the task's time unit. + method: How the records in one step become one record. + + - ``"first"`` / ``"last"``: keep the first / last record, with + its tokens, values and children. Both are over the whole step, + not per token. + - ``"min"`` / ``"max"`` / ``"avg"``: reduce the values. Only for + a level with values and no tokens or children. + - ``"merge"``: union the records' children into one set. Only for + a level with children and no tokens or values, whose child + level is a ``set``. + + A step without records gets ```` tokens, NaN values and no + children; filling it is up to imputation. + """ + + step: float + method: ResampleMethod + + +class LevelConfig(TypedDict, total=False): + """Configuration of one level; a plain ``dict`` with type hints. + + Like :class:`Records`, ``tokens`` and ``values`` declare a role by being + present. They are fixed by the task, so :meth:`SequenceProcessorV2.fit` + only builds vocabularies. + + Keys: + ordering: ``"set"``, ``"sequence"`` (default) or ``"series"``. + Describes the raw input; ``series`` levels carry ``time``. + tokens: Number of token namespaces ``K``; the level's records have + ``tokens``, plain strings when ``K = 1`` and ``K``-tuples otherwise. + values: Value shape ``S``, e.g. ``[1]`` for scalars; the level's + records have ``values`` of shape ``[N, *S]``. + code_mapping: A ``(source, target)`` pair for + ``pyhealth.medcode.CrossMap``, or a list with one pair (or None) + per token namespace. A code that maps to several targets becomes + several records, so a mapping with several targets per code is + only allowed on a ``set`` level without values. + resample: A :class:`Resample` for a ``series`` level: puts its + records on a time grid, turning it into a ``sequence`` level. + spatial: One bool per value axis; defaults to all False. + normalize: Reserved for future work; only None. + impute: Reserved for future work; only None. + missing_indicator: Reserved for future work; only False. + """ + + ordering: OrderType + tokens: int + values: Sequence[int] + code_mapping: CodeMapping | Sequence[CodeMapping | None] | None + resample: Resample | None + spatial: Sequence[bool] | None + normalize: None + impute: None + missing_indicator: Literal[False] + + +@dataclass(frozen=True) +class _InternalLevelConfig: + """A checked :class:`LevelConfig` with every default filled in. + + A level without tokens has ``tokens == 0`` and one without values has + ``values == ()``; neither is a valid setting, since ``K >= 1`` and scalars + are ``(1,)``. + """ + + ordering: OrderType + tokens: int + values: tuple[int, ...] + code_mapping: tuple[CodeMapping | None, ...] | None + resample: tuple[float, ResampleMethod] | None # (step, method) + spatial: tuple[bool, ...] | None + + @staticmethod + def _is_count(value: Any) -> bool: + return isinstance(value, int) and not isinstance(value, bool) and value > 0 + + @staticmethod + def _parse_code_mapping( + level: int, mapping: Any + ) -> tuple[CodeMapping | None, ...] | None: + """Normalizes ``code_mapping`` to one entry (or None) per namespace.""" + if mapping is None: + return None + + def as_pair(entry: Any) -> CodeMapping | None: + if ( + isinstance(entry, (tuple, list)) + and len(entry) == 2 + and all(isinstance(code, str) for code in entry) + ): + return (entry[0], entry[1]) + return None + + pair = as_pair(mapping) + if pair is not None: + return (pair,) + if isinstance(mapping, (tuple, list)) and mapping: + pairs = [None if entry is None else as_pair(entry) for entry in mapping] + if all( + entry is None or pair is not None for entry, pair in zip(mapping, pairs) + ): + return tuple(pairs) + raise ValueError( + f"Level {level}: code_mapping must be a (source, target) pair or a list " + f"with one pair (or None) per token namespace, got {mapping!r}." + ) + + @classmethod + def parse(cls, level: int, config: LevelConfig) -> "_InternalLevelConfig": + """Checks one level's configuration and fills in the defaults. + + ``config`` still comes from user data (a task schema), so every key and + value is checked at runtime despite the annotation. + """ + if not isinstance(config, dict): + raise TypeError( + f"Level {level}: configuration must be a dict, got {type(config).__name__}." + ) + unknown = set(config) - LevelConfig.__annotations__.keys() + if unknown: + raise ValueError( + f"Level {level}: unknown configuration keys {sorted(unknown)}." + ) + + for key, default in _RESERVED.items(): + if config.get(key, default) != default: + raise NotImplementedError( + f"Level {level}: '{key}' is reserved for future work and only " + f"accepts {default!r} for now." + ) + + ordering = config.get("ordering", "sequence") + if ordering not in get_args(OrderType): + raise ValueError( + f"Level {level}: ordering must be one of {get_args(OrderType)}, got {ordering!r}." + ) + + tokens = config.get("tokens", 0) + if "tokens" in config and not cls._is_count(tokens): + raise ValueError( + f"Level {level}: tokens must be a positive number of namespaces, " + f"got {tokens!r}." + ) + + values = config.get("values", ()) + if "values" in config and ( + not isinstance(values, (tuple, list)) + or not values + or not all(cls._is_count(size) for size in values) + ): + raise ValueError( + f"Level {level}: values must be a non-empty list of positive " + f"sizes, e.g. [1] for scalars, got {values!r}." + ) + values = tuple(values) + + if level == 0 and ordering != "series" and not tokens and not values: + raise ValueError( + "Level 0: set 'tokens' or 'values', or use ordering 'series'." + ) + + code_mapping = cls._parse_code_mapping(level, config.get("code_mapping")) + if code_mapping is not None and len(code_mapping) != tokens: + raise ValueError( + f"Level {level}: code_mapping has {len(code_mapping)} entries but the " + f"level has {tokens} token namespace(s)." + ) + resample = config.get("resample") + if resample is not None: + if not isinstance(resample, dict) or resample.keys() != {"step", "method"}: + raise ValueError( + f"Level {level}: resample must be a dict with 'step' and 'method', " + f"got {resample!r}." + ) + step, method = resample["step"], resample["method"] + if ( + isinstance(step, bool) + or not isinstance(step, (int, float)) + or not (np.isfinite(step) and step > 0) + ): + raise ValueError( + f"Level {level}: resample step must be a positive number, got {step!r}." + ) + if method not in get_args(ResampleMethod): + raise ValueError( + f"Level {level}: resample method must be one of " + f"{get_args(ResampleMethod)}, got {method!r}." + ) + if ordering != "series": + raise ValueError( + f"Level {level}: resample requires ordering 'series', got {ordering!r}." + ) + if method in ("min", "max", "avg") and (not values or tokens or level > 0): + raise ValueError( + f"Level {level}: resample method {method!r} reduces values, so the " + "level needs values and no tokens or children (level 0 only)." + ) + if method == "merge" and (level == 0 or tokens or values): + raise ValueError( + f"Level {level}: resample method 'merge' unions children, so the " + "level needs children and no tokens or values." + ) + resample = (float(step), method) + + spatial = config.get("spatial") + if spatial is not None: + if not isinstance(spatial, (tuple, list)) or not all( + isinstance(flag, bool) for flag in spatial + ): + raise ValueError( + f"Level {level}: spatial must be a list of bools, got {spatial!r}." + ) + if len(spatial) != len(values): + raise ValueError( + f"Level {level}: spatial needs one flag per value axis, got " + f"{len(spatial)} flag(s) for values {values}." + ) + spatial = tuple(spatial) + + return cls( + ordering=ordering, + tokens=tokens, + values=values, + code_mapping=code_mapping, + resample=resample, + spatial=spatial, + ) + + +def _new_vocab() -> dict[str, int]: + return {"": TokenProcessorInterface.PAD, "": TokenProcessorInterface.UNK} + + +@register_processor("sequence_v2") +class SequenceProcessorV2(FeatureProcessor): + """Processor for hierarchical packed records, without padding or truncation. + + Reads a :class:`Records` and emits a tuple of tensors in :meth:`schema` + order: by level, then by role (``tokens``, ``values``, ``time``, + ``offsets``). ``offsets`` at level ``l`` partitions level ``l`` records + into level ``l + 1`` records, in CSR layout. Which roles a level has, the + number of token namespaces ``K`` and the value shape ``S`` come from the + configuration; :meth:`fit` only builds the vocabularies. + + Subclasses are configuration only: their ``__init__`` passes fixed + ``levels`` to ``super().__init__()`` and adds no logic. + + Args: + levels: One :class:`LevelConfig` per level, innermost (level 0) + first. + + Examples: + >>> proc = SequenceProcessorV2(levels=[{"ordering": "set", "tokens": 1}, {"ordering": "sequence"}]) + >>> proc.fit([{"visits": {"children": [{"tokens": ["A", "B"]}, {"tokens": ["C"]}]}}], "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])) + """ + + def __init__(self, levels: Sequence[LevelConfig]): + if not isinstance(levels, (list, tuple)) or not levels: + raise ValueError("levels must be a non-empty list of dicts.") + self._config = [ + _InternalLevelConfig.parse(level, config) + for level, config in enumerate(levels) + ] + self.levels = [config.copy() for config in levels] + self._mappers = [ + self._load_mappers(config.code_mapping) for config in self._config + ] + for level, (config, mappers) in enumerate(zip(self._config, self._mappers)): + # A code with several targets becomes several records, which only + # makes sense in an unordered level with no values to copy. + many = [m for m in mappers or () if m is not None and m.has_many_targets] + if many and (config.ordering != "set" or config.values): + raise ValueError( + f"Level {level}: {many[0].s_vocab}->{many[0].t_vocab} maps some " + "codes to several targets, which is only allowed on a 'set' " + f"level without values (got ordering {config.ordering!r}" + f"{' with values' if config.values else ''}). Put the tokens " + "to map in their own 'set' level, one level below this one." + ) + + for level, config in enumerate(self._config): + merge = config.resample is not None and config.resample[1] == "merge" + if merge and self._config[level - 1].ordering != "set": + raise ValueError( + f"Level {level}: resample method 'merge' unions the children's " + f"records, so level {level - 1} must have ordering 'set'." + ) + + # The only fitted state: one vocabulary per level and token namespace. + self._vocabs = [ + [_new_vocab() for _ in range(config.tokens)] for config in self._config + ] + self._fitted = False + + @staticmethod + def _load_mappers( + code_mapping: tuple[CodeMapping | None, ...] | None, + ) -> "tuple[CrossMap | None, ...] | None": + if code_mapping is None: + return None + from pyhealth.medcode import CrossMap + + return tuple( + None if pair is None else CrossMap.load(pair[0], pair[1]) + for pair in code_mapping + ) + + # ------------------------------------------------------------------ + # Metadata + # ------------------------------------------------------------------ + + @property + def depth(self) -> int: + """The outermost level ``L``.""" + return len(self._config) - 1 + + def _require_fit(self) -> None: + if not self._fitted: + raise RuntimeError(f"{type(self).__name__}: fit() must be called first.") + + def _output_roles(self, level: int) -> tuple[Role, ...]: + config = self._config[level] + present = { + "tokens": config.tokens > 0, + "values": len(config.values) > 0, + "time": config.ordering == "series" and config.resample is None, + "offsets": level < self.depth, + } + return tuple(role for role in get_args(Role) if present[role]) + + # FeatureProcessor.schema() returns semantic names for legacy models; V2 + # deliberately returns (level, role) keys instead and is not used by them. + def schema(self) -> tuple[FieldKey, ...]: # type: ignore[override] + """Fields of :meth:`process` output: ``(level, role)``, by level then role.""" + return tuple( + (level, role) + for level in range(self.depth + 1) + for role in self._output_roles(level) + ) + + def batch_schema(self) -> tuple[FieldKey, ...]: + """Fields of :meth:`collate` output: :meth:`schema` plus the batch offsets.""" + batch_offsets: FieldKey = (self.depth, "offsets") + return (*self.schema(), batch_offsets) + + def ordering(self) -> dict[int, OrderType]: + """Ordering of each output level; a resampled level becomes ``sequence``.""" + return { + level: "sequence" if config.resample is not None else config.ordering + for level, config in enumerate(self._config) + } + + def vocab_size(self) -> dict[int, tuple[int, ...] | None]: + """Vocabulary size per token namespace, including ```` and ````.""" + self._require_fit() + return { + level: tuple(len(vocab) for vocab in self._vocabs[level]) + if config.tokens + else None + for level, config in enumerate(self._config) + } + + def value_dim(self) -> dict[int, tuple[int, ...] | None]: + """Trailing value shape ``S`` per level.""" + return { + level: config.values or None for level, config in enumerate(self._config) + } + + # Like schema(): V2 reports one entry per level instead of the legacy tuple. + def spatial(self) -> dict[int, tuple[bool, ...] | None]: # type: ignore[override] + """Whether each value axis is spatial, per level.""" + return { + level: None + if not config.values + else config.spatial or (False,) * len(config.values) + for level, config in enumerate(self._config) + } + + def vocab(self, level: int, namespace: int = 0) -> dict[str, int]: + """Token-to-index mapping of one namespace at one level.""" + self._require_fit() + if not self._config[level].tokens: + raise ValueError(f"Level {level} has no tokens.") + return dict(self._vocabs[level][namespace]) + + # ------------------------------------------------------------------ + # Reading raw records + # ------------------------------------------------------------------ + + def _read_tokens(self, level: int, records: Records) -> list[tuple[str, ...]]: + """Returns each token as a K-tuple, checking it against the configured K.""" + if "tokens" not in records: + raise ValueError(f"Level {level}: the records have no 'tokens'.") + k = self._config[level].tokens + tokens: list[tuple[str, ...]] = [] + for token in records["tokens"]: + if k == 1 and isinstance(token, str): + tokens.append((token,)) + elif ( + k > 1 + and isinstance(token, tuple) + and len(token) == k + and all(isinstance(code, str) for code in token) + ): + tokens.append(token) + else: + expected = "a string" if k == 1 else f"a tuple of {k} strings" + raise ValueError( + f"Level {level}: expected each token as {expected}, got {token!r}." + ) + return tokens + + def _map_tokens( + self, level: int, tokens: list[tuple[str, ...]] + ) -> list[tuple[str, ...]]: + """Applies the code mapping; a code with several targets becomes several records.""" + mappers = self._mappers[level] + if mappers is None: + return tokens + mapped_tokens: list[tuple[str, ...]] = [] + for token in tokens: + targets: list[list[str]] = [] + for code, mapper in zip(token, mappers): + mapped = mapper.map(code) if mapper is not None else [] + targets.append(mapped if mapped else [code]) + if all(len(t) == 1 for t in targets): + mapped_tokens.append(tuple(t[0] for t in targets)) + continue + # The configuration rules out time and values on a mapped level. + if level > 0 or len(token) > 1: + reason = ( + "its records also have children" + if level > 0 + else "there are several namespaces" + ) + raise ValueError( + f"Level {level}: {token!r} maps to several codes, which would duplicate " + f"records, but {reason}." + ) + mapped_tokens.extend((target,) for target in targets[0]) + return mapped_tokens + + def _read_values(self, level: int, records: Records) -> np.ndarray: + """Returns ``float32`` values ``[N, *S]``.""" + if "values" not in records: + raise ValueError(f"Level {level}: the records have no 'values'.") + shape = self._config[level].values + array = np.asarray(records["values"], dtype=np.float64) + if array.ndim == 1 and shape == (1,): + array = array.reshape(-1, 1) # scalars given as [v, ...] + if array.ndim == 1 and array.size == 0: + return np.zeros((0, *shape), dtype=np.float32) + if array.ndim < 2 or tuple(array.shape[1:]) != shape: + raise ValueError( + f"Level {level}: expected values of shape [N, *{shape}], got " + f"{tuple(array.shape)}." + ) + with np.errstate(over="ignore"): # overflow becomes inf, which process rejects + return array.astype(np.float32) + + def _read_time(self, level: int, records: Records) -> np.ndarray: + if "time" not in records: + raise ValueError(f"Level {level}: the 'series' records have no 'time'.") + array = np.asarray(records["time"], dtype=np.float32) + if array.ndim != 1: + raise ValueError( + f"Level {level}: time must be 1-D, got shape {array.shape}." + ) + if not np.isfinite(array).all(): + raise ValueError(f"Level {level}: time must be finite.") + if (np.diff(array) < 0).any(): + raise ValueError(f"Level {level}: time must be non-decreasing.") + return array + + def _resample( + self, + level: int, + step: float, + method: ResampleMethod, + columns: dict[Role, np.ndarray], + children: list[list[Records]], + ) -> tuple[dict[Role, np.ndarray], list[list[Records]]]: + """Puts one level's records on a time grid, one record per step. + + Args: + columns: The level's ``tokens`` (IDs), ``values`` and ``time`` for + all its records, and ``offsets``: the record count of each + group of records (one group per record of the level above). + children: One group of child ``Records`` per record. + + Returns: + The columns and children of the grid: one record per step, no + ``time``, and ``offsets`` holding each group's number of steps. A + step without records gets ```` tokens, NaN values and no + children, for imputation to fill. + """ + time, counts = columns["time"], columns["offsets"] + if (time < 0).any(): + raise ValueError(f"Level {level}: resample needs non-negative time.") + steps = np.floor(time.astype(np.float64) / step).astype(np.int64) + + # A group's grid runs from step 0 to its last record's step. Each group + # gets its own block of the level's grid, so `grid` never decreases. + sizes = np.zeros(len(counts), dtype=np.int64) + filled_groups = counts > 0 + sizes[filled_groups] = steps[np.cumsum(counts)[filled_groups] - 1] + 1 + grid = steps + np.repeat(np.cumsum(sizes) - sizes, counts) + n = int(sizes.sum()) + + # The first and last record of each step that has records. + starts = np.flatnonzero(np.diff(grid, prepend=-1)) + ends = np.flatnonzero(np.diff(grid, append=-1)) + filled = grid[starts] + + resampled: dict[Role, np.ndarray] = {"offsets": sizes} + grid_children: list[list[Records]] = [[] for _ in range(n if level else 0)] + if method == "merge": + for i, kids in zip(grid.tolist(), children): + grid_children[i].extend(kids) + elif method == "first" or method == "last": + keep = starts if method == "first" else ends + pad = TokenProcessorInterface.PAD + for role, fill in (("tokens", pad), ("values", np.nan)): + if role in columns: + column = columns[role] + on_grid = np.full((n, *column.shape[1:]), fill, dtype=column.dtype) + on_grid[filled] = column[keep] + resampled[role] = on_grid + if level: + for i, kept in zip(filled.tolist(), keep.tolist()): + grid_children[i] = children[kept] + else: + reduce = {"min": np.minimum, "max": np.maximum, "avg": np.add}[method] + values = reduce.reduceat( + columns["values"].astype(np.float64), starts, axis=0 + ) + if method == "avg": + values /= (ends - starts + 1).reshape(-1, *[1] * (values.ndim - 1)) + on_grid = np.full((n, *values.shape[1:]), np.nan, dtype=np.float32) + on_grid[filled] = values + resampled["values"] = on_grid + return resampled, grid_children + + def _empty(self, level: int, role: Role) -> np.ndarray: + """The start of each column, so a level without records still works.""" + config = self._config[level] + if role == "tokens": + return np.zeros((0, config.tokens), dtype=np.int64) + if role == "values": + return np.zeros((0, *config.values), dtype=np.float32) + return np.zeros(0, dtype=np.float32) # time + + # ------------------------------------------------------------------ + # Pipeline + # ------------------------------------------------------------------ + + def fit(self, samples: Iterable[dict[str, Any]], field: str) -> None: + """Builds the vocabularies from the (code-mapped) tokens. + + Records are not checked here; :meth:`process` checks them. + + Args: + samples: Iterable of sample dicts, read once. + field: Key whose values are :class:`Records`. + """ + for sample in samples: + # Walk one level at a time, from the outermost level L down to 0. + level_records: list[Records] = [sample[field]] + for level in range(self.depth, -1, -1): + vocabs, mappers = self._vocabs[level], self._mappers[level] + children: list[Records] = [] + for records in level_records: + for token in records.get("tokens", ()): + codes = (token,) if isinstance(token, str) else token + for vocab, code, mapper in zip( + vocabs, codes, mappers or repeat(None) + ): + targets = mapper.map(code) if mapper is not None else [] + for target in targets or [code]: + if target not in vocab: + vocab[target] = len(vocab) + children.extend(records.get("children", ())) + level_records = children + self._fitted = True + + def process(self, value: Records) -> tuple[torch.Tensor, ...]: + """Turns one :class:`Records` into packed tensors in :meth:`schema` order. + + Args: + value: The sample's :class:`Records` for this field. + + Returns: + Tuple of tensors: ``int64`` tokens (``[N]``, or ``[N, K]``), + ``float32`` values ``[N, *S]``, ``float32`` time ``[N]`` and + ``int64`` offsets ``[N_{l+1} + 1]``. + """ + self._require_fit() + unk = TokenProcessorInterface.UNK + fields: dict[FieldKey, torch.Tensor] = {} + # Walk one level at a time, from the outermost level L down to 0. Each + # record of a level owns one group of Records one level down; the + # groups come in their records' order, which is CSR order. + groups: list[list[Records]] = [[value]] + for level in range(self.depth, -1, -1): + config, vocabs = self._config[level], self._vocabs[level] + declared: tuple[tuple[Role, bool], ...] = ( + ("tokens", config.tokens > 0), + ("values", len(config.values) > 0), + ("time", config.ordering == "series"), + ) + parts: dict[Role, list[np.ndarray]] = { + role: [self._empty(level, role)] for role, on in declared if on + } + counts: list[int] = [] + children: list[list[Records]] = [] + for group in groups: + count = 0 + for records in group: + # Only what the configuration declares is read; other keys + # are ignored. + lengths: dict[str, int] = {} + tokens = None + if config.tokens: + tokens = self._read_tokens(level, records) + lengths["tokens"] = len(tokens) + if config.values: + values = self._read_values(level, records) + parts["values"].append(values) + lengths["values"] = len(values) + if config.ordering == "series": + time = self._read_time(level, records) + parts["time"].append(time) + lengths["time"] = len(time) + if level > 0: + if "children" not in records: + raise ValueError( + f"Level {level}: the records have no 'children' " + f"({self.depth + 1} levels are configured)." + ) + children.extend([kid] for kid in records["children"]) + lengths["children"] = len(records["children"]) + if len(set(lengths.values())) > 1: + raise ValueError( + f"Level {level}: columns have different lengths {lengths}." + ) + n = next(iter(lengths.values())) + if tokens is not None: + tokens = self._map_tokens(level, tokens) + ids = [ + [vocab.get(code, unk) for vocab, code in zip(vocabs, token)] + for token in tokens + ] + parts["tokens"].append( + np.array(ids, dtype=np.int64).reshape( + len(tokens), config.tokens + ) + ) + n = len(tokens) + count += n + counts.append(count) + + columns: dict[Role, np.ndarray] = { + role: np.concatenate(part) for role, part in parts.items() + } + columns["offsets"] = np.array(counts, dtype=np.int64) + if config.resample is not None: + columns, children = self._resample( + level, *config.resample, columns, children + ) + for role in self._output_roles(level): + array = columns[role] + if role == "offsets": + array = np.concatenate([np.zeros(1, np.int64), np.cumsum(array)]) + elif role == "tokens" and config.tokens == 1: + array = array.reshape(-1) + elif role == "values" and not np.isfinite(array).all(): + raise ValueError( + f"Level {level}: values contain NaN or inf, and imputation " + "is not supported yet." + ) + fields[(level, role)] = torch.from_numpy(array) + groups = children + return tuple(fields[key] for key in self.schema()) + + def collate( + self, items: list[tuple[torch.Tensor, ...]] + ) -> tuple[torch.Tensor, ...]: + """Packs a batch of :meth:`process` outputs, in :meth:`batch_schema` order. + + Payloads are concatenated, each sample's offsets are shifted by the + records before it, and the batch offsets ``[B + 1]`` are appended. + """ + if not items: + raise ValueError("collate needs at least one sample.") + schema = self.schema() + output = [] + for index, (_, role) in enumerate(schema): + column = [item[index] for item in items] + if role != "offsets": + output.append(torch.cat(column)) + continue + parts, shift = [column[0][:1]], 0 + for offsets in column: + parts.append(offsets[1:] + shift) + shift += int(offsets[-1]) + output.append(torch.cat(parts)) + counts = [self._count(schema, item, self.depth) for item in items] + output.append(torch.tensor([0, *accumulate(counts)], dtype=torch.long)) + return tuple(output) + + @staticmethod + def _count( + schema: tuple[FieldKey, ...], output: tuple[torch.Tensor, ...], level: int + ) -> int: + """Number of records at ``level``, read from any field that knows it.""" + for (field_level, role), tensor in zip(schema, output): + if field_level == level and role != "offsets": + return tensor.shape[0] + if field_level == level - 1 and role == "offsets": + return tensor.shape[0] - 1 + raise ValueError(f"No field records the size of level {level}.") + + def validate(self, output: tuple[torch.Tensor, ...], batched: bool = False) -> None: + """Checks an output of :meth:`process` (or :meth:`collate` if ``batched``). + + Raises: + TypeError: If a field is not a tensor. + ValueError: If the output breaks the packed-record invariants. + """ + self._require_fit() + schema = self.batch_schema() if batched else self.schema() + if not isinstance(output, tuple) or len(output) != len(schema): + raise ValueError(f"Expected a tuple of {len(schema)} tensors for {schema}.") + fields = dict(zip(schema, output)) + + for (level, role), tensor in fields.items(): + if not isinstance(tensor, torch.Tensor): + raise TypeError( + f"{(level, role)}: expected a tensor, got {type(tensor).__name__}." + ) + dtype = torch.long if role in ("tokens", "offsets") else torch.float32 + if tensor.dtype != dtype: + raise ValueError( + f"{(level, role)}: expected {dtype}, got {tensor.dtype}." + ) + if role == "tokens": + sizes = tuple(len(vocab) for vocab in self._vocabs[level]) + expected_ndim = 1 if len(sizes) == 1 else 2 + if tensor.dim() != expected_ndim or ( + expected_ndim == 2 and tensor.shape[1] != len(sizes) + ): + raise ValueError( + f"{(level, role)}: bad shape {tuple(tensor.shape)} for {len(sizes)} namespace(s)." + ) + limits = ( + torch.tensor(sizes) + .expand(tensor.shape[0], -1) + .reshape(tensor.shape) + ) + if (tensor < 0).any() or (tensor >= limits).any(): + raise ValueError( + f"{(level, role)}: token index outside the vocabulary." + ) + elif role == "values": + shape = self._config[level].values + if tuple(tensor.shape[1:]) != shape or tensor.dim() != 1 + len(shape): + raise ValueError( + f"{(level, role)}: expected [N, *{shape}], got {tuple(tensor.shape)}." + ) + if not torch.isfinite(tensor).all(): + raise ValueError(f"{(level, role)}: values contain NaN or inf.") + elif role == "time": + if tensor.dim() != 1 or not torch.isfinite(tensor).all(): + raise ValueError(f"{(level, role)}: time must be 1-D and finite.") + else: + if ( + tensor.dim() != 1 + or tensor.numel() == 0 + or tensor[0] != 0 + or (tensor.diff() < 0).any() + ): + raise ValueError( + f"{(level, role)}: offsets must start at 0 and never decrease." + ) + + for level in range(self.depth + 1): + sizes = { + (field_level, role): tensor.shape[0] + for (field_level, role), tensor in fields.items() + if field_level == level and role != "offsets" + } + if level > 0: + sizes[(level - 1, "offsets")] = ( + fields[(level - 1, "offsets")].shape[0] - 1 + ) + if len(set(sizes.values())) > 1: + raise ValueError( + f"Level {level}: fields disagree on the record count {sizes}." + ) + offsets = fields.get((level, "offsets")) + if ( + offsets is not None + and sizes + and offsets[-1] != next(iter(sizes.values())) + ): + raise ValueError( + f"{(level, 'offsets')}: last offset must equal the level's record count." + ) + + def __repr__(self) -> str: + return f"{type(self).__name__}(levels={self.levels})" diff --git a/tests/core/test_sequence_processor_v2.py b/tests/core/test_sequence_processor_v2.py new file mode 100644 index 000000000..cbd114492 --- /dev/null +++ b/tests/core/test_sequence_processor_v2.py @@ -0,0 +1,714 @@ +import pickle +import unittest +from collections.abc import Sequence +from typing import Any +from unittest.mock import patch + +import torch + +from pyhealth.processors import ( + LevelConfig, + Records, + Resample, + SequenceProcessorV2, + get_processor, +) + +NAN = float("nan") +CODES: LevelConfig = {"tokens": 1} +SCALARS: LevelConfig = {"values": [1]} +NESTED: list[LevelConfig] = [{"ordering": "set", "tokens": 1}, {"ordering": "sequence"}] +LAST: Resample = {"step": 1.0, "method": "last"} + + +def _fit(levels: Sequence[LevelConfig], *records: Any) -> SequenceProcessorV2: + proc = SequenceProcessorV2(levels=levels) + proc.fit([{"x": r} for r in records], "x") + return proc + + +def _visits(*visits: list[str]) -> Records: + return {"children": [{"tokens": list(v)} for v in visits]} + + +class _FakeMap: + def __init__(self, table): + self.table = table + self.s_vocab, self.t_vocab = "ICD9CM", "CCSCM" + self.has_many_targets = any(len(t) > 1 for t in table.values()) + + def map(self, code): + return self.table.get(code, []) + + +class TestConfiguration(unittest.TestCase): + def test_levels_are_required(self): + with self.assertRaises(TypeError): + SequenceProcessorV2() # type: ignore[call-arg] # the missing argument is the test + + def test_registered(self): + self.assertIs(get_processor("sequence_v2"), SequenceProcessorV2) + + def test_subclass_is_configuration_only(self): + class NestedCodes(SequenceProcessorV2): + def __init__(self): + super().__init__(levels=NESTED) + + proc = NestedCodes() + proc.fit([{"x": _visits(["A"])}], "x") + self.assertEqual(proc.schema(), ((0, "tokens"), (0, "offsets"))) + + def test_rejects_bad_configurations(self): + mapping = ("ICD9CM", "CCSCM") + # Deliberately invalid, so typed Any rather than LevelConfig. + cases: list[Any] = [ + [], + ["set"], + [{"order": "set", "tokens": 1}], + [{"ordering": "bag", "tokens": 1}], + [{}], # level 0 declares nothing + [{"tokens": 0}], + [{"tokens": True}], + [{"tokens": 1.0}], + [{"values": []}], + [{"values": 1}], + [{"values": [0]}], + [{"values": [True]}], + [{"values": [1], "resample": LAST}], # resample on a sequence level + [{"ordering": "series", "values": [1], "resample": 1.0}], # not a dict + [{"ordering": "series", "values": [1], "resample": {"step": 1.0}}], + [{"ordering": "series", "values": [1], "resample": {**LAST, "step": 0}}], + [{"ordering": "series", "values": [1], "resample": {**LAST, "step": True}}], + [ + { + "ordering": "series", + "values": [1], + "resample": {**LAST, "method": "sum"}, + } + ], + [{"values": [1, 1], "spatial": [1, 0]}], + [{"tokens": 1, "spatial": [False]}], # spatial without values + [{"values": [1], "spatial": [False, True]}], # one flag per axis + [{"tokens": 1, "code_mapping": "ICD9CM"}], + [{"tokens": 2, "code_mapping": [mapping, "x"]}], + [{"tokens": 2, "code_mapping": mapping}], # one entry per namespace + [{"values": [1], "code_mapping": mapping}], # mapping without tokens + ] + for levels in cases: + with ( + self.subTest(levels=levels), + self.assertRaises((ValueError, TypeError)), + ): + SequenceProcessorV2(levels=levels) + + def test_reserved_options_not_implemented(self): + reserved: list[tuple[str, Any]] = [ + ("normalize", "zscore"), + ("impute", "zero"), + ("missing_indicator", True), + ] + for key, value in reserved: + config: Any = {"tokens": 1, key: value} # deliberately invalid + with self.subTest(key=key), self.assertRaises(NotImplementedError): + SequenceProcessorV2(levels=[config]) + SequenceProcessorV2( + levels=[ + { + "tokens": 1, + "normalize": None, + "impute": None, + "missing_indicator": False, + } + ] + ) + + def test_only_vocabulary_needs_fit(self): + proc = SequenceProcessorV2(levels=[{"tokens": 2, "values": [3]}]) + self.assertEqual(proc.schema(), ((0, "tokens"), (0, "values"))) + self.assertEqual(proc.value_dim(), {0: (3,)}) + self.assertEqual(proc.spatial(), {0: (False,)}) + with self.assertRaises(RuntimeError): + proc.vocab_size() + with self.assertRaises(RuntimeError): + proc.vocab(0) + with self.assertRaises(RuntimeError): + proc.process({"tokens": [("A", "u")], "values": [[1.0, 2.0, 3.0]]}) + + +class TestSpecExamples(unittest.TestCase): + """The examples of Section 5.1, written as Records.""" + + def assertOutput(self, output, expected): + self.assertEqual(len(output), len(expected)) + for got, want in zip(output, expected): + self.assertEqual(got.dtype, want.dtype) + torch.testing.assert_close(got, want) + + def test_flat_sequence(self): + raw: Records = {"tokens": ["D01", "D02", "D01"]} + proc = _fit([CODES], raw) + self.assertEqual(proc.schema(), ((0, "tokens"),)) + self.assertOutput(proc.process(raw), (torch.tensor([2, 3, 2]),)) + self.assertEqual(proc.vocab_size(), {0: (4,)}) + self.assertEqual(proc.value_dim(), {0: None}) + self.assertEqual(proc.spatial(), {0: None}) + self.assertEqual(proc.ordering(), {0: "sequence"}) + + def test_codes_in_visits(self): + raw = _visits(["D01", "D02"], [], ["D03"]) + proc = _fit(NESTED, raw) + self.assertEqual(proc.schema(), ((0, "tokens"), (0, "offsets"))) + self.assertOutput( + proc.process(raw), (torch.tensor([2, 3, 4]), torch.tensor([0, 2, 2, 3])) + ) + self.assertEqual(proc.vocab_size(), {0: (5,), 1: None}) + self.assertEqual(proc.ordering(), {0: "set", 1: "sequence"}) + + def test_admissions_visits_codes(self): + raw: Records = {"children": [_visits(["A", "B"], ["C"]), _visits(["A", "D"])]} + proc = _fit([{"ordering": "set", "tokens": 1}, {}, {}], raw) + self.assertEqual(proc.schema(), ((0, "tokens"), (0, "offsets"), (1, "offsets"))) + self.assertOutput( + proc.process(raw), + ( + torch.tensor([2, 3, 4, 2, 5]), + torch.tensor([0, 2, 3, 5]), + torch.tensor([0, 2, 3]), + ), + ) + self.assertEqual(proc.vocab_size(), {0: (6,), 1: None, 2: None}) + + def test_codes_values_and_time(self): + raw: Records = { + "time": [0.0, 1.5, 6.0], + "tokens": ["Glucose", "WBC", "Glucose"], + "values": [[105.0], [7.2], [110.0]], + } + proc = _fit([{"ordering": "series", "tokens": 1, "values": [1]}], raw) + self.assertEqual(proc.schema(), ((0, "tokens"), (0, "values"), (0, "time"))) + self.assertOutput( + proc.process(raw), + ( + torch.tensor([2, 3, 2]), + torch.tensor([[105.0], [7.2], [110.0]]), + torch.tensor([0.0, 1.5, 6.0]), + ), + ) + self.assertEqual(proc.value_dim(), {0: (1,)}) + self.assertEqual(proc.spatial(), {0: (False,)}) + self.assertEqual(proc.ordering(), {0: "series"}) + + def test_multi_namespace_tokens(self): + raw: Records = { + "tokens": [("Glucose", "mg/dL"), ("WBC", "10^3/uL")], + "values": [[105.0], [7.2]], + } + proc = _fit([{"tokens": 2, "values": [1]}], raw) + self.assertEqual(proc.schema(), ((0, "tokens"), (0, "values"))) + self.assertOutput( + proc.process(raw), + (torch.tensor([[2, 2], [3, 3]]), torch.tensor([[105.0], [7.2]])), + ) + self.assertEqual(proc.vocab_size(), {0: (4, 4)}) + self.assertEqual( + proc.vocab(0, 1), {"": 0, "": 1, "mg/dL": 2, "10^3/uL": 3} + ) + + def test_multichannel_time_series(self): + raw: Records = { + "time": [0.0, 1.0, 2.0], + "values": [[80.0, 95.0], [82.0, 91.0], [78.0, 94.0]], + } + proc = _fit([{"ordering": "series", "values": [2]}], raw) + self.assertEqual(proc.schema(), ((0, "values"), (0, "time"))) + output = proc.process(raw) + self.assertEqual(output[0].shape, (3, 2)) + self.assertEqual(proc.vocab_size(), {0: None}) + self.assertEqual(proc.value_dim(), {0: (2,)}) + + def test_notes_over_time(self): + raw: Records = { + "time": [0.0, 3.5], + "children": [ + {"tokens": ["[CLS]", "pain", "[SEP]"]}, + {"tokens": ["[CLS]", "[SEP]"]}, + ], + } + proc = _fit([CODES, {"ordering": "series"}], raw) + self.assertEqual(proc.schema(), ((0, "tokens"), (0, "offsets"), (1, "time"))) + self.assertOutput( + proc.process(raw), + ( + torch.tensor([2, 3, 4, 2, 4]), + torch.tensor([0, 3, 5]), + torch.tensor([0.0, 3.5]), + ), + ) + self.assertEqual(proc.ordering(), {0: "sequence", 1: "series"}) + + def test_spatial_values(self): + raw: Records = {"time": [0.0, 24.0], "values": torch.zeros(2, 1, 4, 4).numpy()} + config: LevelConfig = { + "ordering": "series", + "values": [1, 4, 4], + "spatial": [False, True, True], + } + proc = _fit([config], raw) + self.assertEqual(proc.process(raw)[0].shape, (2, 1, 4, 4)) + self.assertEqual(proc.value_dim(), {0: (1, 4, 4)}) + self.assertEqual(proc.spatial(), {0: (False, True, True)}) + + +class TestProcess(unittest.TestCase): + def test_unseen_tokens_are_unknown_and_nothing_is_truncated(self): + proc = _fit(NESTED, _visits(["A"])) + tokens, offsets = proc.process(_visits(["A", "B", "C", "D"], ["A"])) + self.assertEqual(tokens.tolist(), [2, 1, 1, 1, 2]) + self.assertEqual(offsets.tolist(), [0, 4, 5]) + + def test_fit_grows_vocabulary_across_samples(self): + proc = _fit([CODES], {"tokens": ["A"]}, {"tokens": ["B", "A"]}) + self.assertEqual(proc.vocab(0), {"": 0, "": 1, "A": 2, "B": 3}) + + def test_empty_inputs(self): + proc = _fit([{"tokens": 2}], {"tokens": [("a", "b")]}) + self.assertEqual(proc.process({"tokens": []})[0].shape, (0, 2)) + + proc = _fit( + [{"ordering": "series", "values": [2]}], + {"time": [0.0], "values": [[1.0, 2.0]]}, + ) + values, time = proc.process({"time": [], "values": []}) + self.assertEqual(values.shape, (0, 2)) + self.assertEqual(time.shape, (0,)) + + proc = _fit(NESTED, _visits(["A"])) + tokens, offsets = proc.process({"children": []}) + self.assertEqual(tokens.shape, (0,)) + self.assertEqual(offsets.tolist(), [0]) + + def test_fit_needs_no_records(self): + # Shapes come from the configuration, so empty training data is fine. + proc = _fit(NESTED, {"children": []}) + self.assertEqual(proc.vocab_size(), {0: (2,), 1: None}) + self.assertEqual(proc.process(_visits(["A"]))[0].tolist(), [1]) + + proc = _fit([SCALARS], {"values": []}) + self.assertEqual(proc.process({"values": [[1.0]]})[0].shape, (1, 1)) + + def test_dtypes(self): + raw: Records = { + "time": [0, 1], + "tokens": ["A", "B"], + "values": [[1], [2]], + } + tokens, values, time = _fit( + [{"ordering": "series", "tokens": 1, "values": [1]}], raw + ).process(raw) + self.assertEqual( + (tokens.dtype, values.dtype, time.dtype), + (torch.long, torch.float32, torch.float32), + ) + + def test_flat_scalars_become_one_column(self): + proc = _fit([SCALARS], {"values": [[1.0]]}) + (values,) = proc.process({"values": [1.0, 2.0]}) + self.assertEqual(values.tolist(), [[1.0], [2.0]]) + with self.assertRaises(ValueError): # only when S is (1,) + _fit([{"values": [2]}], {"values": [[1.0, 2.0]]}).process( + {"values": [1.0, 2.0]} + ) + + def test_undeclared_keys_are_ignored(self): + proc = _fit([CODES], {"tokens": ["A"]}) + # Typed Any: "extra" is deliberately not a Records key. + cases: tuple[Any, ...] = ( + {"tokens": ["A"], "values": [[1.0]]}, # values not declared + {"tokens": ["A"], "extra": [1]}, + {"tokens": ["A"], "time": [0.0]}, # time on a sequence level + {"tokens": ["A"], "children": [{"tokens": ["B"]}]}, # level 0 + ) + for raw in cases: + with self.subTest(raw=raw): + self.assertEqual(proc.process(raw)[0].tolist(), [2]) + + def test_bad_records_raise_when_read(self): + series: LevelConfig = {"ordering": "series", "tokens": 1} + pairs: LevelConfig = {"tokens": 2} + cases: list[tuple[list[LevelConfig], Any]] = [ + ([CODES], {"values": [[1.0]]}), # declared tokens missing + ([CODES], {"tokens": [("A", "u")]}), # K = 1 takes strings + ([CODES], {"tokens": [("A",)]}), + ([pairs], {"tokens": ["A"]}), + ([pairs], {"tokens": [("A", "u", "v")]}), + ([CODES], {"tokens": [None]}), + ( + [{"tokens": 1, "values": [1]}], + {"tokens": ["A"], "values": [[1.0], [2.0]]}, + ), # lengths differ + ([CODES], ["A"]), + ([CODES], {}), + ([series], {"tokens": ["A"]}), # series without time + ([series], {"time": [1.0, 0.0], "tokens": ["A", "B"]}), + ([series], {"time": [NAN], "tokens": ["A"]}), + ([SCALARS], {"values": [[1.0, 2.0]]}), + ([SCALARS], {"values": [[[1.0]]]}), + ([SCALARS], {"values": [[NAN]]}), + (NESTED, {"tokens": ["A"]}), # no children at level 1 + ] + for levels, raw in cases: + proc = _fit(levels) + with ( + self.subTest(raw=raw), + self.assertRaises((ValueError, TypeError)), + ): + proc.process(raw) + + def test_process_checks_records(self): + proc = _fit( + [{"tokens": 1, "values": [1]}], {"tokens": ["A"], "values": [[1.0]]} + ) + cases: tuple[Records, ...] = ( + {"tokens": ["A"]}, + {"tokens": ["A"], "values": [[1.0, 2.0]]}, + {"tokens": ["A"], "values": [[NAN]]}, + {"tokens": ["A"], "values": [[float("inf")]]}, + {"tokens": ["A"], "values": [[1e39]]}, # overflows float32 + ) + for raw in cases: + with self.subTest(raw=raw), self.assertRaises(ValueError): + proc.process(raw) + + +class TestCodeMapping(unittest.TestCase): + MAPPING = ("ICD9CM", "CCSCM") + + def setUp(self): + self.use_table({"250.00": ["49"], "401.9": ["98", "99"]}) # multi-target + + def use_table(self, table): + patcher = patch("pyhealth.medcode.CrossMap.load", return_value=_FakeMap(table)) + self.load = patcher.start() + self.addCleanup(patcher.stop) + + def test_maps_and_keeps_unmapped_codes(self): + proc = _fit( + [{"ordering": "set", "tokens": 1, "code_mapping": self.MAPPING}], + {"tokens": ["250.00", "X"]}, + ) + self.load.assert_called_once_with("ICD9CM", "CCSCM") + self.assertEqual(proc.vocab(0), {"": 0, "": 1, "49": 2, "X": 3}) + self.assertEqual( + proc.process({"tokens": ["250.00", "X", "Y"]})[0].tolist(), [2, 3, 1] + ) + + def test_one_to_many_expands_records(self): + levels: list[LevelConfig] = [ + {"ordering": "set", "tokens": 1, "code_mapping": self.MAPPING}, + {}, + ] + raw = _visits(["401.9", "250.00"], ["401.9"]) + tokens, offsets = _fit(levels, raw).process(raw) + self.assertEqual(tokens.tolist(), [2, 3, 4, 2, 3]) + self.assertEqual(offsets.tolist(), [0, 3, 5]) + + def test_per_namespace_mapping(self): + proc = _fit( + [{"ordering": "set", "tokens": 2, "code_mapping": [None, ("A", "B")]}], + {"tokens": [("401.9", "250.00")]}, + ) + self.assertEqual(proc.vocab(0, 0), {"": 0, "": 1, "401.9": 2}) + self.assertEqual(proc.vocab(0, 1), {"": 0, "": 1, "49": 2}) + + def test_multi_target_mapping_needs_a_set_level_without_values(self): + mapping = self.MAPPING + cases: list[list[LevelConfig]] = [ + [{"tokens": 1, "code_mapping": mapping}], # sequence + [{"ordering": "series", "tokens": 1, "code_mapping": mapping}], + [{"ordering": "set", "tokens": 1, "values": [1], "code_mapping": mapping}], + ] + for levels in cases: + with self.subTest(levels=levels), self.assertRaises(ValueError): + SequenceProcessorV2(levels=levels) + + def test_single_target_mapping_is_allowed_anywhere(self): + self.use_table({"250.00": ["49"], "401.9": ["98"]}) + raw: Records = { + "time": [0.0, 1.0], + "tokens": ["250.00", "X"], + "values": [[1.0], [2.0]], + } + levels: list[LevelConfig] = [ + { + "ordering": "series", + "tokens": 1, + "values": [1], + "code_mapping": self.MAPPING, + } + ] + tokens, values, time = _fit(levels, raw).process(raw) + self.assertEqual(tokens.tolist(), [2, 3]) + self.assertEqual(values.shape, (2, 1)) + self.assertEqual(time.tolist(), [0.0, 1.0]) + + def test_expansion_raises_when_it_would_duplicate_data(self): + mapping = self.MAPPING + cases: list[tuple[list[LevelConfig], Records]] = [ + ( + [{"ordering": "set", "tokens": 2, "code_mapping": [mapping, None]}], + {"tokens": [("401.9", "u")]}, + ), + ( + [CODES, {"ordering": "set", "tokens": 1, "code_mapping": mapping}], + {"tokens": ["401.9"], "children": [{"tokens": ["A"]}]}, + ), + ] + for levels, raw in cases: + proc = _fit(levels, raw) # fit only builds vocabularies + with self.subTest(raw=raw), self.assertRaises(ValueError): + proc.process(raw) + + +class TestResample(unittest.TestCase): + @staticmethod + def series(method: Any, **config: Any) -> LevelConfig: + # Typed Any: the tests also build invalid configurations. + level: Any = {"ordering": "series", "resample": {"step": 1.0, "method": method}} + level.update(config) + return level + + def test_methods_need_the_right_roles(self): + series = self.series + cases: list[list[LevelConfig]] = [ + [series("avg", tokens=1, values=[1])], # reduces values only + [series("min")], # no values + [CODES, series("max", values=[1])], # level 1 has children + [series("merge", tokens=1)], # level 0 has no children + [CODES, series("merge", values=[1])], + [CODES, series("merge", tokens=1)], + [{"tokens": 1}, series("merge")], # child level is not a set + ] + for levels in cases: + with self.subTest(levels=levels), self.assertRaises(ValueError): + SequenceProcessorV2(levels=levels) + + def test_first_and_last_keep_one_record_per_step(self): + raw: Records = { + "time": [0.0, 0.5, 1.2, 2.0], + "values": [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]], + } + for method, expected in ( + ("first", [[1.0, 2.0], [5.0, 6.0], [7.0, 8.0]]), + ("last", [[3.0, 4.0], [5.0, 6.0], [7.0, 8.0]]), + ): + with self.subTest(method=method): + proc = _fit([self.series(method, values=[2])], raw) + self.assertEqual(proc.schema(), ((0, "values"),)) + self.assertEqual(proc.ordering(), {0: "sequence"}) + (values,) = proc.process(raw) + self.assertEqual(values.tolist(), expected) + + def test_last_keeps_tokens_and_children(self): + raw: Records = { + "time": [0.0, 0.5, 1.5], + "tokens": ["a", "b", "c"], + "children": [{"tokens": ["X"]}, {"tokens": ["Y", "Z"]}, {"tokens": []}], + } + proc = _fit([CODES, self.series("last", tokens=1)], raw) + self.assertEqual(proc.schema(), ((0, "tokens"), (0, "offsets"), (1, "tokens"))) + child_tokens, offsets, tokens = proc.process(raw) + self.assertEqual(tokens.tolist(), [3, 4]) # "b", "c" + self.assertEqual(child_tokens.tolist(), [3, 4]) # "Y", "Z" + self.assertEqual(offsets.tolist(), [0, 2, 2]) + proc.validate((child_tokens, offsets, tokens)) + + def test_steps_without_records_are_left_for_imputation(self): + raw: Records = { + "time": [0.0, 2.0], + "tokens": ["a", "b"], + "children": [{"tokens": ["X"]}, {"tokens": ["Y"]}], + } + proc = _fit([CODES, self.series("first", tokens=1)], raw) + child_tokens, offsets, tokens = proc.process(raw) + self.assertEqual(tokens.tolist(), [2, 0, 3]) # step 1 is + self.assertEqual(child_tokens.tolist(), [2, 3]) + self.assertEqual(offsets.tolist(), [0, 1, 1, 2]) # and has no children + proc.validate((child_tokens, offsets, tokens)) + + def test_each_group_gets_its_own_grid(self): + stay_a: Records = {"time": [0.0, 1.5], "values": [[1.0], [2.0]]} + stay_b: Records = {"time": [], "values": []} + stay_c: Records = {"time": [0.2, 0.7], "values": [[3.0], [4.0]]} + raw: Records = {"children": [stay_a, stay_b, stay_c]} + proc = _fit([self.series("last", values=[1]), {}], raw) + values, offsets = proc.process(raw) + self.assertEqual(values.tolist(), [[1.0], [2.0], [4.0]]) + self.assertEqual(offsets.tolist(), [0, 2, 2, 3]) + + def test_min_max_avg_reduce_values(self): + raw: Records = { + "time": [0.0, 0.5, 1.0], + "values": [[1.0, 6.0], [3.0, 2.0], [5.0, 5.0]], + } + for method, expected in ( + ("min", [[1.0, 2.0], [5.0, 5.0]]), + ("max", [[3.0, 6.0], [5.0, 5.0]]), + ("avg", [[2.0, 4.0], [5.0, 5.0]]), + ): + with self.subTest(method=method): + proc = _fit([self.series(method, values=[2])], raw) + self.assertEqual(proc.process(raw)[0].tolist(), expected) + + def test_merge_unions_child_sets(self): + raw: Records = { + "time": [0.0, 0.5, 2.0], + "children": [{"tokens": ["A"]}, {"tokens": ["B", "C"]}, {"tokens": ["D"]}], + } + proc = _fit([{"ordering": "set", "tokens": 1}, self.series("merge")], raw) + self.assertEqual(proc.schema(), ((0, "tokens"), (0, "offsets"))) + tokens, offsets = proc.process(raw) + self.assertEqual(tokens.tolist(), [2, 3, 4, 5]) + self.assertEqual(offsets.tolist(), [0, 3, 3, 4]) # step 1 is an empty set + proc.validate((tokens, offsets)) + + def test_merge_regroups_below_the_top_level(self): + days: Records = { + "time": [0.0, 0.5], + "children": [{"tokens": ["A"]}, {"tokens": ["B"]}], + } + raw: Records = {"children": [days, days]} + levels = [{"ordering": "set", "tokens": 1}, self.series("merge"), {}] + proc = _fit(levels, raw) + tokens, day_offsets, stay_offsets = proc.process(raw) + self.assertEqual(tokens.tolist(), [2, 3, 2, 3]) + self.assertEqual(day_offsets.tolist(), [0, 2, 4]) # one merged day per stay + self.assertEqual(stay_offsets.tolist(), [0, 1, 2]) + batch = proc.collate([(tokens, day_offsets, stay_offsets)] * 2) + proc.validate(batch, batched=True) + + def test_grid_starts_at_time_zero(self): + proc = _fit([self.series("last", values=[1])]) + # Step 0 has no record, so its values are NaN, which imputation (not + # supported yet) would have to fill. + with self.assertRaisesRegex(ValueError, "imputation"): + proc.process({"time": [1.0, 2.0], "values": [[1.0], [2.0]]}) + self.assertEqual( + proc.process({"time": [0.0, 1.0], "values": [[1.0], [2.0]]})[0].shape, + (2, 1), + ) + + def test_value_gaps_and_negative_time_raise(self): + proc = _fit([self.series("avg", values=[1])]) + cases: tuple[Records, ...] = ( + {"time": [0.0, 2.0], "values": [[1.0], [2.0]]}, + {"time": [-1.0, 0.0], "values": [[1.0], [2.0]]}, + ) + for raw in cases: + with self.subTest(raw=raw), self.assertRaises(ValueError): + proc.process(raw) + + def test_empty_series(self): + proc = _fit([self.series("last", values=[1])]) + self.assertEqual(proc.process({"time": [], "values": []})[0].shape, (0, 1)) + + def test_records_need_values(self): + with self.assertRaises(ValueError): + _fit([self.series("last", values=[1])]).process({"time": [0.0]}) + + +class TestCollateAndValidate(unittest.TestCase): + def test_collate_nested(self): + a, b = _visits(["A", "B"], ["C"]), _visits(["C"], [], ["A"]) + proc = _fit(NESTED, a, b) + batch = proc.collate( + [proc.process(a), proc.process(b), proc.process({"children": []})] + ) + self.assertEqual( + proc.batch_schema(), ((0, "tokens"), (0, "offsets"), (1, "offsets")) + ) + self.assertEqual(batch[0].tolist(), [2, 3, 4, 4, 2]) + self.assertEqual(batch[1].tolist(), [0, 2, 3, 4, 4, 5]) + self.assertEqual(batch[2].tolist(), [0, 2, 5, 5]) + proc.validate(batch, batched=True) + + def test_collate_deep_and_flat_series(self): + deep: Records = {"children": [_visits(["A"], ["B", "C"]), _visits(["D"])]} + proc = _fit([{"ordering": "set", "tokens": 1}, {}, {}], deep) + batch = proc.collate([proc.process(deep)] * 2) + self.assertEqual(batch[1].tolist(), [0, 1, 3, 4, 5, 7, 8]) + self.assertEqual(batch[2].tolist(), [0, 2, 3, 5, 6]) + self.assertEqual(batch[3].tolist(), [0, 2, 4]) + proc.validate(batch, batched=True) + + raw: Records = {"time": [0.0, 1.0], "values": [[1.0], [2.0]]} + proc = _fit([{"ordering": "series", "values": [1]}], raw) + values, time, batch_offsets = proc.collate( + [proc.process(raw), proc.process({"time": [5.0], "values": [[3.0]]})] + ) + self.assertEqual(values.shape, (3, 1)) + self.assertEqual(time.tolist(), [0.0, 1.0, 5.0]) + self.assertEqual(batch_offsets.tolist(), [0, 2, 3]) + + def test_collate_needs_items(self): + with self.assertRaises(ValueError): + _fit([CODES], {"tokens": ["A"]}).collate([]) + + def test_validate_accepts_process_output(self): + raw: Records = { + "time": [0.0, 1.0], + "children": [ + {"tokens": [("A", "u")], "values": [[1.0]]}, + {"tokens": [], "values": []}, + ], + } + proc = _fit([{"tokens": 2, "values": [1]}, {"ordering": "series"}], raw) + proc.validate(proc.process(raw)) + + def test_validate_rejects_broken_outputs(self): + proc = _fit(NESTED, _visits(["A", "B"], ["C"])) + tokens, offsets = proc.process(_visits(["A", "B"], ["C"])) + cases = [ + (tokens,), + [tokens, offsets], + (tokens.int(), offsets), + (torch.tensor([2, 3, 9]), offsets), + (tokens, torch.tensor([1, 2, 3])), + (tokens, torch.tensor([0, 3, 2])), + (tokens, torch.tensor([0, 2, 4])), + (tokens[:2], offsets), + ] + for output in cases: + with self.subTest(output=output), self.assertRaises(ValueError): + proc.validate(output) + with self.assertRaises(ValueError): + proc.validate((tokens, offsets, torch.tensor([0, 1])), batched=True) + + +class TestStorage(unittest.TestCase): + def test_pickle_round_trip(self): + raw = _visits(["A", "B"], ["C"]) + proc = pickle.loads(pickle.dumps(_fit(NESTED, raw))) + self.assertEqual(proc.process(raw)[0].tolist(), [2, 3, 4]) + + def test_sample_dataset_round_trip(self): + from pyhealth.datasets import create_sample_dataset + + samples = [ + {"patient_id": "p0", "visits": _visits(["A", "B"], ["C"]), "label": 0}, + {"patient_id": "p1", "visits": _visits(["C"]), "label": 1}, + ] + dataset = create_sample_dataset( + samples=samples, + input_schema={"visits": ("sequence_v2", {"levels": NESTED})}, + output_schema={"label": "binary"}, + ) + proc = dataset.input_processors["visits"] + for sample, raw in zip(dataset, samples): + expected = proc.process(raw["visits"]) + self.assertIsInstance(sample["visits"], tuple) + for got, want in zip(sample["visits"], expected): + torch.testing.assert_close(got, want) + + +if __name__ == "__main__": + unittest.main() From 000f427f34254199a24cadef2ca6f826180e59da Mon Sep 17 00:00:00 2001 From: Logiquo Date: Wed, 28 Oct 2026 08:00:06 -0400 Subject: [PATCH 03/11] Add imputation to SequenceProcessorV2 and split out resampling - Impute: a step {"method", "by_tokens", "by_dims"} or a chain of them, run after resampling. Methods: zero (0.0 / ), mean / min / max (dataset statistics of the raw values, learned in fit(), per token in by_tokens and reduced over the value axes in by_dims), forward_fill / backward_fill / linear (along time, within one group, optionally only from records with the same token). Gaps a step cannot fill keep NaN and still raise unless a later step fills them. - learns_statistics is True when an impute step uses statistics. - Resampling and imputation now live in private _Resampler and _Imputer classes; fit() and process() only call them. --- ...yhealth.processors.SequenceProcessorV2.rst | 10 +- pyhealth/processors/__init__.py | 2 + pyhealth/processors/sequence_processor_v2.py | 550 +++++++++++++++--- tests/core/test_sequence_processor_v2.py | 166 +++++- 4 files changed, 635 insertions(+), 93 deletions(-) diff --git a/docs/api/processors/pyhealth.processors.SequenceProcessorV2.rst b/docs/api/processors/pyhealth.processors.SequenceProcessorV2.rst index 0389b5588..a064c5b0f 100644 --- a/docs/api/processors/pyhealth.processors.SequenceProcessorV2.rst +++ b/docs/api/processors/pyhealth.processors.SequenceProcessorV2.rst @@ -35,9 +35,10 @@ In a task schema the configuration is plain data: } Batches are packed with ``collate()``, which follows ``batch_schema()``: the -fields of ``schema()`` plus the batch offsets ``[B + 1]``. ``normalize``, -``impute`` and ``missing_indicator`` are reserved for future work and only -accept ``None`` / ``False``; values containing NaN or inf raise. +fields of ``schema()`` plus the batch offsets ``[B + 1]``. ``impute`` fills NaN +values (see ``Impute``); ``normalize`` and ``missing_indicator`` are reserved +for future work and only accept ``None`` / ``False``. Values still containing +NaN or inf raise. .. autoclass:: pyhealth.processors.SequenceProcessorV2 :members: @@ -52,3 +53,6 @@ accept ``None`` / ``False``; values containing NaN or inf raise. .. autoclass:: pyhealth.processors.Resample :members: + +.. autoclass:: pyhealth.processors.Impute + :members: diff --git a/pyhealth/processors/__init__.py b/pyhealth/processors/__init__.py index dcf368bee..28b3a15fe 100644 --- a/pyhealth/processors/__init__.py +++ b/pyhealth/processors/__init__.py @@ -38,6 +38,7 @@ def get_processor(name: str): from .raw_processor import RawProcessor from .sequence_processor import SequenceProcessor from .sequence_processor_v2 import ( + Impute, LevelConfig, Records, Resample, @@ -85,6 +86,7 @@ def get_processor(name: str): "SequenceProcessorV2", "Records", "LevelConfig", + "Impute", "Resample", "SignalProcessor", "StageNetProcessor", diff --git a/pyhealth/processors/sequence_processor_v2.py b/pyhealth/processors/sequence_processor_v2.py index 96c64b2d7..d421ef869 100644 --- a/pyhealth/processors/sequence_processor_v2.py +++ b/pyhealth/processors/sequence_processor_v2.py @@ -3,10 +3,18 @@ See ``notes/Unified_Sequence_Processor.md`` (Section 7) for the design. """ -from collections.abc import Iterable, Sequence +from collections.abc import Iterable, Iterator, Sequence from dataclasses import dataclass from itertools import accumulate, repeat -from typing import TYPE_CHECKING, Any, Literal, TypedDict, get_args +from typing import ( + TYPE_CHECKING, + Any, + Literal, + NamedTuple, + Required, + TypedDict, + get_args, +) import numpy as np import numpy.typing as npt @@ -23,10 +31,17 @@ FieldKey = tuple[int, Role] CodeMapping = tuple[str, str] ResampleMethod = Literal["first", "last", "min", "max", "avg", "merge"] +ImputeMethod = Literal[ + "zero", "mean", "min", "max", "forward_fill", "backward_fill", "linear" +] +# Imputation methods that use statistics learned in fit(). +_STATISTICS = ("mean", "min", "max") +# Imputation methods that work along time. +_ALONG_TIME = ("forward_fill", "backward_fill", "linear") # Options reserved in the configuration but not implemented in Phase 1. -_RESERVED = {"normalize": None, "impute": None, "missing_indicator": False} +_RESERVED = {"normalize": None, "missing_indicator": False} class Records(TypedDict, total=False): @@ -83,12 +98,53 @@ class Resample(TypedDict): method: ResampleMethod +class Impute(TypedDict, total=False): + """One imputation step for a level's NaN values; a plain ``dict``. + + Keys: + method: How to fill. + + - ``"zero"``: 0.0 for values and ```` for tokens. Ignores + ``by_tokens`` and ``by_dims``. + - ``"mean"`` / ``"min"`` / ``"max"``: a statistic of the raw + training values (before resampling), learned in ``fit()`` + with NaN skipped. + - ``"forward_fill"`` / ``"backward_fill"``: the previous / next + value along time. + - ``"linear"``: linear interpolation along time, using ``time`` + (or the step index on a resampled level). + + The last three need a ``sequence`` or ``series`` level, never + cross into another group of records, work value position by + value position (ignoring ``by_dims``), and leave gaps at the + edges as NaN. + by_tokens: Token namespaces. Statistics are kept per token in + these namespaces, and the time-based methods only use records + with the same token. None (the default) ignores tokens; the time- + based methods then also fill ```` tokens. + by_dims: For statistics: the axes of the value shape ``S`` to + reduce over, indexed like ``values`` and ``spatial``. None (the + default) reduces over all of them, ``()`` over none (one + statistic per value position); e.g. for ``values: [3, 64, 64]``, + ``(1, 2)`` gives one statistic per channel. Time is always + reduced. + + A key without a statistic (a token or value never seen in ``fit()``) + keeps its NaN, as does any gap a method cannot fill; chain several steps + to fill those. + """ + + method: Required[ImputeMethod] + by_tokens: Sequence[int] | None + by_dims: Sequence[int] | None + + class LevelConfig(TypedDict, total=False): """Configuration of one level; a plain ``dict`` with type hints. Like :class:`Records`, ``tokens`` and ``values`` declare a role by being present. They are fixed by the task, so :meth:`SequenceProcessorV2.fit` - only builds vocabularies. + only builds vocabularies (and impute statistics). Keys: ordering: ``"set"``, ``"sequence"`` (default) or ``"series"``. @@ -105,8 +161,9 @@ class LevelConfig(TypedDict, total=False): resample: A :class:`Resample` for a ``series`` level: puts its records on a time grid, turning it into a ``sequence`` level. spatial: One bool per value axis; defaults to all False. + impute: An :class:`Impute` step, or a list of them applied in order, + run after resampling. normalize: Reserved for future work; only None. - impute: Reserved for future work; only None. missing_indicator: Reserved for future work; only False. """ @@ -116,11 +173,19 @@ class LevelConfig(TypedDict, total=False): code_mapping: CodeMapping | Sequence[CodeMapping | None] | None resample: Resample | None spatial: Sequence[bool] | None + impute: Impute | Sequence[Impute] | None normalize: None - impute: None missing_indicator: Literal[False] +class _ImputeStep(NamedTuple): + """A checked :class:`Impute`.""" + + method: ImputeMethod + by_tokens: tuple[int, ...] | None + by_dims: tuple[int, ...] | None + + @dataclass(frozen=True) class _InternalLevelConfig: """A checked :class:`LevelConfig` with every default filled in. @@ -136,11 +201,18 @@ class _InternalLevelConfig: code_mapping: tuple[CodeMapping | None, ...] | None resample: tuple[float, ResampleMethod] | None # (step, method) spatial: tuple[bool, ...] | None + impute: tuple[_ImputeStep, ...] @staticmethod def _is_count(value: Any) -> bool: return isinstance(value, int) and not isinstance(value, bool) and value > 0 + @staticmethod + def _is_index(value: Any, size: int) -> bool: + return ( + isinstance(value, int) and not isinstance(value, bool) and 0 <= value < size + ) + @staticmethod def _parse_code_mapping( level: int, mapping: Any @@ -284,6 +356,65 @@ def parse(cls, level: int, config: LevelConfig) -> "_InternalLevelConfig": ) spatial = tuple(spatial) + impute = config.get("impute") + steps = ( + [] if impute is None else [impute] if isinstance(impute, dict) else impute + ) + if not isinstance(steps, (list, tuple)): + raise TypeError( + f"Level {level}: impute must be a dict or a list of dicts, got {impute!r}." + ) + parsed: list[_ImputeStep] = [] + for step in steps: + if ( + not isinstance(step, dict) + or "method" not in step + or not step.keys() <= Impute.__annotations__.keys() + ): + raise ValueError( + f"Level {level}: each impute step must be a dict with 'method' and " + f"optionally 'by_tokens' and 'by_dims', got {step!r}." + ) + method = step["method"] + if method not in get_args(ImputeMethod): + raise ValueError( + f"Level {level}: impute method must be one of " + f"{get_args(ImputeMethod)}, got {method!r}." + ) + by_tokens, by_dims = step.get("by_tokens"), step.get("by_dims") + if by_tokens is not None: + if ( + not isinstance(by_tokens, (tuple, list)) + or not by_tokens + or not all(cls._is_index(k, tokens) for k in by_tokens) + ): + raise ValueError( + f"Level {level}: by_tokens must be None or a non-empty list of " + f"token namespaces below {tokens}, got {by_tokens!r}." + ) + by_tokens = tuple(by_tokens) + if by_dims is not None: + if ( + not isinstance(by_dims, (tuple, list)) + or not all(cls._is_index(d, len(values)) for d in by_dims) + or len(set(by_dims)) != len(by_dims) + ): + raise ValueError( + f"Level {level}: by_dims must be None or a list of distinct " + f"value axes below {len(values)}, got {by_dims!r}." + ) + by_dims = tuple(by_dims) + if method in (*_STATISTICS, "linear") and not values: + raise ValueError( + f"Level {level}: impute method {method!r} needs values." + ) + if method in _ALONG_TIME and ordering == "set": + raise ValueError( + f"Level {level}: impute method {method!r} works along time, so the " + "level must be a 'sequence' or 'series'." + ) + parsed.append(_ImputeStep(method, by_tokens, by_dims)) + return cls( ordering=ordering, tokens=tokens, @@ -291,9 +422,298 @@ def parse(cls, level: int, config: LevelConfig) -> "_InternalLevelConfig": code_mapping=code_mapping, resample=resample, spatial=spatial, + impute=tuple(parsed), ) +class _Resampler: + """Puts the records of one ``series`` level on a time grid; see :class:`Resample`.""" + + def __init__(self, level: int, step: float, method: ResampleMethod): + self.level, self.step, self.method = level, step, method + + def __call__( + self, columns: dict[Role, np.ndarray], children: list[list[Records]] + ) -> tuple[dict[Role, np.ndarray], list[list[Records]]]: + """Puts the level's records on the grid, one record per step. + + Args: + columns: The level's ``tokens`` (IDs), ``values`` and ``time`` for + all its records, and ``offsets``: the record count of each + group of records (one group per record of the level above). + children: One group of child ``Records`` per record. + + Returns: + The columns and children of the grid: one record per step, no + ``time``, and ``offsets`` holding each group's number of steps. A + step without records gets ```` tokens, NaN values and no + children, for imputation to fill. + """ + time, counts = columns["time"], columns["offsets"] + if (time < 0).any(): + raise ValueError(f"Level {self.level}: resample needs non-negative time.") + steps = np.floor(time.astype(np.float64) / self.step).astype(np.int64) + + # A group's grid runs from step 0 to its last record's step. Each group + # gets its own block of the level's grid, so `grid` never decreases. + sizes = np.zeros(len(counts), dtype=np.int64) + filled_groups = counts > 0 + sizes[filled_groups] = steps[np.cumsum(counts)[filled_groups] - 1] + 1 + grid = steps + np.repeat(np.cumsum(sizes) - sizes, counts) + n = int(sizes.sum()) + + # The first and last record of each step that has records. + starts = np.flatnonzero(np.diff(grid, prepend=-1)) + ends = np.flatnonzero(np.diff(grid, append=-1)) + filled = grid[starts] + + resampled: dict[Role, np.ndarray] = {"offsets": sizes} + grid_children: list[list[Records]] = [[] for _ in range(n if self.level else 0)] + if self.method == "merge": + for i, kids in zip(grid.tolist(), children): + grid_children[i].extend(kids) + elif self.method == "first" or self.method == "last": + keep = starts if self.method == "first" else ends + pad = TokenProcessorInterface.PAD + for role, fill in (("tokens", pad), ("values", np.nan)): + if role in columns: + column = columns[role] + on_grid = np.full((n, *column.shape[1:]), fill, dtype=column.dtype) + on_grid[filled] = column[keep] + resampled[role] = on_grid + if self.level: + for i, kept in zip(filled.tolist(), keep.tolist()): + grid_children[i] = children[kept] + else: + reduce = {"min": np.minimum, "max": np.maximum, "avg": np.add}[self.method] + values = reduce.reduceat( + columns["values"].astype(np.float64), starts, axis=0 + ) + if self.method == "avg": + values /= (ends - starts + 1).reshape(-1, *[1] * (values.ndim - 1)) + on_grid = np.full((n, *values.shape[1:]), np.nan, dtype=np.float32) + on_grid[filled] = values + resampled["values"] = on_grid + return resampled, grid_children + + +class _Imputer: + """Runs the impute steps of one level (see :class:`Impute`). + + It also holds the statistics of the ``mean`` / ``min`` / ``max`` steps: + :meth:`accumulate` adds one ``Records``' raw values during ``fit()``, and + :meth:`finish` turns the running state into statistics. + """ + + def __init__(self, level: int, steps: tuple[_ImputeStep, ...], tokens: int): + self.level, self.steps, self.tokens = level, steps, tokens + self.learns_statistics = any(step.method in _STATISTICS for step in steps) + # Per statistic step: its running state during fit(), then, per token + # key, its statistic. + self._state = self._new_state() + self._stats: dict[int, dict[tuple[int, ...], np.ndarray]] = {} + + def _new_state(self) -> dict[int, dict[tuple[int, ...], list[np.ndarray]]]: + return { + index: {} + for index, step in enumerate(self.steps) + if step.method in _STATISTICS + } + + def accumulate(self, values: np.ndarray, ids: list[list[int]]) -> None: + """Adds one ``Records``' raw values (and their token IDs) to the statistics.""" + for index, state in self._state.items(): + step = self.steps[index] + keys = np.zeros((len(values), 0), dtype=np.int64) + if step.by_tokens is not None: + if len(ids) != len(values): + raise ValueError( + f"Level {self.level}: {len(ids)} tokens but {len(values)} values." + ) + table = np.array(ids, dtype=np.int64).reshape(len(ids), self.tokens) + keys = table[:, list(step.by_tokens)] + self._accumulate(step, values, keys, state) + + def finish(self) -> None: + """Turns the running state into statistics, ready for a new fit().""" + for index, state in self._state.items(): + stats: dict[tuple[int, ...], np.ndarray] = {} + for key, parts in state.items(): + if self.steps[index].method == "mean": + with np.errstate(invalid="ignore", divide="ignore"): + stat = parts[0] / parts[1] + else: + stat = np.where(np.isinf(parts[0]), np.nan, parts[0]) + stats[key] = stat.astype(np.float32) + self._stats[index] = stats + self._state = self._new_state() + + @staticmethod + def _split_by_key(keys: np.ndarray) -> Iterator[tuple[tuple[int, ...], np.ndarray]]: + """Yields each distinct row of ``keys`` and the mask of rows equal to it.""" + if len(keys) == 0: + return + if keys.shape[1] == 0: + yield (), np.ones(len(keys), dtype=bool) + return + unique, inverse = np.unique(keys, axis=0, return_inverse=True) + inverse = inverse.reshape(-1) + for index, key in enumerate(unique.tolist()): + yield tuple(key), inverse == index + + def _accumulate( + self, + step: _ImputeStep, + values: np.ndarray, + keys: np.ndarray, + state: dict[tuple[int, ...], list[np.ndarray]], + ) -> None: + """Adds raw values to the running state of one statistic, per token key.""" + reduced = range(values.ndim - 1) if step.by_dims is None else step.by_dims + axes = (0, *(1 + d for d in reduced)) + values = values.astype(np.float64) # sums over a dataset need the range + finite = np.isfinite(values) + for key, rows in self._split_by_key(keys): + x, ok = values[rows], finite[rows] + if step.method == "mean": + total = np.where(ok, x, 0.0).sum(axis=axes, keepdims=True)[0] + count = ok.sum(axis=axes, keepdims=True)[0] + if key in state: + state[key][0] += total + state[key][1] += count + else: + state[key] = [total, count] + else: + low = step.method == "min" + fill = np.inf if low else -np.inf + reduce = np.min if low else np.max + extreme = reduce(np.where(ok, x, fill), axis=axes, keepdims=True)[0] + if key in state: + combine = np.minimum if low else np.maximum + state[key][0] = combine(state[key][0], extreme) + else: + state[key] = [extreme] + + @staticmethod + def _previous_valid(valid: np.ndarray, start: np.ndarray) -> np.ndarray: + """Index of the latest valid row at or before each row, in its segment. + + ``valid`` is ``[L, M]``; ``start[i]`` is the first row of row ``i``'s + segment. Returns ``[L, M]``, with -1 where there is none. + """ + index = np.arange(len(valid))[:, None] + latest = np.maximum.accumulate(np.where(valid, index, -1), axis=0) + return np.where(latest >= start[:, None], latest, -1) + + def _next_valid( + self, valid: np.ndarray, start: np.ndarray, end: np.ndarray + ) -> np.ndarray: + """Like :meth:`_previous_valid`, for the earliest valid row at or after.""" + last = len(valid) - 1 + found = self._previous_valid(valid[::-1], (last - end)[::-1])[::-1] + return np.where(found >= 0, last - found, -1) + + @staticmethod + def _segments( + group: np.ndarray, keys: np.ndarray + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Orders rows by segment: by group, then by token key, keeping time order. + + Returns the order and, for each row in that order, the first and last + row of its segment. + """ + n = len(group) + order = np.lexsort([np.arange(n), *keys.T[::-1], group]) + group, keys = group[order], keys[order] + new = np.ones(n, dtype=bool) + new[1:] = (group[1:] != group[:-1]) | (keys[1:] != keys[:-1]).any(axis=1) + first = np.flatnonzero(new) + segment = np.cumsum(new) - 1 + return order, first[segment], np.append(first[1:], n)[segment] - 1 + + def __call__(self, columns: dict[Role, np.ndarray]) -> None: + """Runs the impute steps in place on the level's concatenated columns.""" + pad = TokenProcessorInterface.PAD + counts = columns["offsets"] + n = int(counts.sum()) + group = np.repeat(np.arange(len(counts)), counts) + for index, step in enumerate(self.steps): + values, tokens = columns.get("values"), columns.get("tokens") + if step.method == "zero": + if values is not None: + columns["values"] = np.where(np.isnan(values), 0.0, values).astype( + np.float32 + ) + continue # missing tokens are already + + keys = ( + tokens[:, list(step.by_tokens)] + if tokens is not None and step.by_tokens is not None + else np.zeros((n, 0), dtype=np.int64) + ) + if step.method in _STATISTICS: + assert values is not None # the configuration requires values + stats = self._stats.get(index, {}) + values = values.copy() + for key, rows in self._split_by_key(keys): + stat = stats.get(key) + if stat is not None: + x = values[rows] + values[rows] = np.where(np.isnan(x), stat, x) + columns["values"] = values + continue + + # Along time: work in segment order, so each segment is contiguous. + order, start, end = self._segments(group, keys) + skip = (keys[order] == pad).any(axis=1) # rows without a token + if values is not None: + x = values.reshape(n, -1)[order] + valid = ~np.isnan(x) + before = self._previous_valid(valid, start) + after = self._next_valid(valid, start, end) + if step.method == "linear": + time = columns.get("time") + position = (time if time is not None else np.arange(n))[order] + position = position.astype(np.float64) + ok = ~valid & (before >= 0) & (after >= 0) + x0 = np.take_along_axis(x, np.maximum(before, 0), axis=0) + x1 = np.take_along_axis(x, np.maximum(after, 0), axis=0) + t0 = position[np.maximum(before, 0)] + t1 = position[np.maximum(after, 0)] + with np.errstate(invalid="ignore", divide="ignore"): + weight = np.where( + t1 > t0, (position[:, None] - t0) / (t1 - t0), 0.0 + ) + filled = np.where(ok, x0 + (x1 - x0) * weight, x) + else: + source = before if step.method == "forward_fill" else after + taken = np.take_along_axis(x, np.maximum(source, 0), axis=0) + filled = np.where(~valid & (source >= 0), taken, x) + filled[skip] = x[skip] + result = np.empty_like(filled) + result[order] = filled + columns["values"] = result.reshape(values.shape).astype(np.float32) + if ( + tokens is not None + and step.by_tokens is None + and step.method != "linear" + ): + # Without by_tokens, tokens (steps without records) are + # filled from the same neighbour as the values. + t = tokens[order] + valid = (t != pad).any(axis=1)[:, None] + if step.method == "forward_fill": + source = self._previous_valid(valid, start)[:, 0] + else: + source = self._next_valid(valid, start, end)[:, 0] + t = np.where( + (~valid[:, 0] & (source >= 0))[:, None], t[np.maximum(source, 0)], t + ) + result = np.empty_like(t) + result[order] = t + columns["tokens"] = result + + def _new_vocab() -> dict[str, int]: return {"": TokenProcessorInterface.PAD, "": TokenProcessorInterface.UNK} @@ -307,7 +727,8 @@ class SequenceProcessorV2(FeatureProcessor): ``offsets``). ``offsets`` at level ``l`` partitions level ``l`` records into level ``l + 1`` records, in CSR layout. Which roles a level has, the number of token namespaces ``K`` and the value shape ``S`` come from the - configuration; :meth:`fit` only builds the vocabularies. + configuration; :meth:`fit` only builds the vocabularies and the statistics + of the impute steps. Subclasses are configuration only: their ``__init__`` passes fixed ``levels`` to ``super().__init__()`` and adds no logic. @@ -357,7 +778,20 @@ def __init__(self, levels: Sequence[LevelConfig]): f"records, so level {level - 1} must have ordering 'set'." ) - # The only fitted state: one vocabulary per level and token namespace. + self._resamplers = [ + None if config.resample is None else _Resampler(level, *config.resample) + for level, config in enumerate(self._config) + ] + # Fitted state: one vocabulary per level and token namespace, and the + # statistics held by each level's imputer. + self._imputers = [ + _Imputer(level, config.impute, config.tokens) if config.impute else None + for level, config in enumerate(self._config) + ] + self.learns_statistics = any( + imputer is not None and imputer.learns_statistics + for imputer in self._imputers + ) self._vocabs = [ [_new_vocab() for _ in range(config.tokens)] for config in self._config ] @@ -543,75 +977,6 @@ def _read_time(self, level: int, records: Records) -> np.ndarray: raise ValueError(f"Level {level}: time must be non-decreasing.") return array - def _resample( - self, - level: int, - step: float, - method: ResampleMethod, - columns: dict[Role, np.ndarray], - children: list[list[Records]], - ) -> tuple[dict[Role, np.ndarray], list[list[Records]]]: - """Puts one level's records on a time grid, one record per step. - - Args: - columns: The level's ``tokens`` (IDs), ``values`` and ``time`` for - all its records, and ``offsets``: the record count of each - group of records (one group per record of the level above). - children: One group of child ``Records`` per record. - - Returns: - The columns and children of the grid: one record per step, no - ``time``, and ``offsets`` holding each group's number of steps. A - step without records gets ```` tokens, NaN values and no - children, for imputation to fill. - """ - time, counts = columns["time"], columns["offsets"] - if (time < 0).any(): - raise ValueError(f"Level {level}: resample needs non-negative time.") - steps = np.floor(time.astype(np.float64) / step).astype(np.int64) - - # A group's grid runs from step 0 to its last record's step. Each group - # gets its own block of the level's grid, so `grid` never decreases. - sizes = np.zeros(len(counts), dtype=np.int64) - filled_groups = counts > 0 - sizes[filled_groups] = steps[np.cumsum(counts)[filled_groups] - 1] + 1 - grid = steps + np.repeat(np.cumsum(sizes) - sizes, counts) - n = int(sizes.sum()) - - # The first and last record of each step that has records. - starts = np.flatnonzero(np.diff(grid, prepend=-1)) - ends = np.flatnonzero(np.diff(grid, append=-1)) - filled = grid[starts] - - resampled: dict[Role, np.ndarray] = {"offsets": sizes} - grid_children: list[list[Records]] = [[] for _ in range(n if level else 0)] - if method == "merge": - for i, kids in zip(grid.tolist(), children): - grid_children[i].extend(kids) - elif method == "first" or method == "last": - keep = starts if method == "first" else ends - pad = TokenProcessorInterface.PAD - for role, fill in (("tokens", pad), ("values", np.nan)): - if role in columns: - column = columns[role] - on_grid = np.full((n, *column.shape[1:]), fill, dtype=column.dtype) - on_grid[filled] = column[keep] - resampled[role] = on_grid - if level: - for i, kept in zip(filled.tolist(), keep.tolist()): - grid_children[i] = children[kept] - else: - reduce = {"min": np.minimum, "max": np.maximum, "avg": np.add}[method] - values = reduce.reduceat( - columns["values"].astype(np.float64), starts, axis=0 - ) - if method == "avg": - values /= (ends - starts + 1).reshape(-1, *[1] * (values.ndim - 1)) - on_grid = np.full((n, *values.shape[1:]), np.nan, dtype=np.float32) - on_grid[filled] = values - resampled["values"] = on_grid - return resampled, grid_children - def _empty(self, level: int, role: Role) -> np.ndarray: """The start of each column, so a level without records still works.""" config = self._config[level] @@ -626,9 +991,10 @@ def _empty(self, level: int, role: Role) -> np.ndarray: # ------------------------------------------------------------------ def fit(self, samples: Iterable[dict[str, Any]], field: str) -> None: - """Builds the vocabularies from the (code-mapped) tokens. + """Builds the vocabularies and the statistics of the impute steps. - Records are not checked here; :meth:`process` checks them. + Vocabularies come from the (code-mapped) tokens, statistics from the + raw values. Records are not checked here; :meth:`process` checks them. Args: samples: Iterable of sample dicts, read once. @@ -639,10 +1005,13 @@ def fit(self, samples: Iterable[dict[str, Any]], field: str) -> None: level_records: list[Records] = [sample[field]] for level in range(self.depth, -1, -1): vocabs, mappers = self._vocabs[level], self._mappers[level] + imputer = self._imputers[level] children: list[Records] = [] for records in level_records: + ids: list[list[int]] = [] for token in records.get("tokens", ()): codes = (token,) if isinstance(token, str) else token + row: list[int] = [] for vocab, code, mapper in zip( vocabs, codes, mappers or repeat(None) ): @@ -650,8 +1019,18 @@ def fit(self, samples: Iterable[dict[str, Any]], field: str) -> None: for target in targets or [code]: if target not in vocab: vocab[target] = len(vocab) + row.append(vocab[(targets or [code])[0]]) + ids.append(row) + if imputer is not None and imputer.learns_statistics: + # Statistics are taken from the raw values, before + # resampling. + imputer.accumulate(self._read_values(level, records), ids) children.extend(records.get("children", ())) level_records = children + + for imputer in self._imputers: + if imputer is not None: + imputer.finish() self._fitted = True def process(self, value: Records) -> tuple[torch.Tensor, ...]: @@ -734,10 +1113,11 @@ def process(self, value: Records) -> tuple[torch.Tensor, ...]: role: np.concatenate(part) for role, part in parts.items() } columns["offsets"] = np.array(counts, dtype=np.int64) - if config.resample is not None: - columns, children = self._resample( - level, *config.resample, columns, children - ) + resampler, imputer = self._resamplers[level], self._imputers[level] + if resampler is not None: + columns, children = resampler(columns, children) + if imputer is not None: + imputer(columns) for role in self._output_roles(level): array = columns[role] if role == "offsets": @@ -746,8 +1126,8 @@ def process(self, value: Records) -> tuple[torch.Tensor, ...]: array = array.reshape(-1) elif role == "values" and not np.isfinite(array).all(): raise ValueError( - f"Level {level}: values contain NaN or inf, and imputation " - "is not supported yet." + f"Level {level}: values contain NaN or inf; impute them " + "(see Impute) or remove them in the task." ) fields[(level, role)] = torch.from_numpy(array) groups = children diff --git a/tests/core/test_sequence_processor_v2.py b/tests/core/test_sequence_processor_v2.py index cbd114492..e12ae5726 100644 --- a/tests/core/test_sequence_processor_v2.py +++ b/tests/core/test_sequence_processor_v2.py @@ -7,6 +7,7 @@ import torch from pyhealth.processors import ( + Impute, LevelConfig, Records, Resample, @@ -104,7 +105,6 @@ def test_rejects_bad_configurations(self): def test_reserved_options_not_implemented(self): reserved: list[tuple[str, Any]] = [ ("normalize", "zscore"), - ("impute", "zero"), ("missing_indicator", True), ] for key, value in reserved: @@ -116,7 +116,6 @@ def test_reserved_options_not_implemented(self): { "tokens": 1, "normalize": None, - "impute": None, "missing_indicator": False, } ] @@ -588,9 +587,8 @@ def test_merge_regroups_below_the_top_level(self): def test_grid_starts_at_time_zero(self): proc = _fit([self.series("last", values=[1])]) - # Step 0 has no record, so its values are NaN, which imputation (not - # supported yet) would have to fill. - with self.assertRaisesRegex(ValueError, "imputation"): + # Step 0 has no record, so its values are NaN, and nothing imputes them. + with self.assertRaisesRegex(ValueError, "NaN"): proc.process({"time": [1.0, 2.0], "values": [[1.0], [2.0]]}) self.assertEqual( proc.process({"time": [0.0, 1.0], "values": [[1.0], [2.0]]})[0].shape, @@ -616,6 +614,164 @@ def test_records_need_values(self): _fit([self.series("last", values=[1])]).process({"time": [0.0]}) +class TestImpute(unittest.TestCase): + @staticmethod + def level(*steps: Any, **config: Any) -> LevelConfig: + # Typed Any: the tests also build invalid configurations. + level: Any = {"impute": list(steps)} + level.update(config) + return level + + def test_rejects_bad_steps(self): + level = self.level + bad: list[Any] = [ + [{"values": [1], "impute": "zero"}], + [level({"by_tokens": [0]}, values=[1])], # no method + [level({"method": "median"}, values=[1])], + [level({"method": "zero", "extra": 1}, values=[1])], + [level({"method": "mean", "by_tokens": [1]}, tokens=1, values=[1])], + [level({"method": "mean", "by_tokens": []}, tokens=1, values=[1])], + [level({"method": "mean", "by_dims": [1]}, values=[1])], + [level({"method": "mean", "by_dims": [0, 0]}, values=[2, 2])], + [level({"method": "mean"}, tokens=1)], # no values + [level({"method": "linear"}, ordering="series", tokens=1)], + [level({"method": "forward_fill"}, ordering="set", values=[1])], + ] + for levels in bad: + with ( + self.subTest(levels=levels), + self.assertRaises((ValueError, TypeError)), + ): + SequenceProcessorV2(levels=levels) + + def test_learns_statistics_only_with_statistics(self): + zero = SequenceProcessorV2(levels=[self.level({"method": "zero"}, values=[1])]) + mean = SequenceProcessorV2(levels=[self.level({"method": "mean"}, values=[1])]) + self.assertFalse(zero.learns_statistics) + self.assertTrue(mean.learns_statistics) + + def test_zero(self): + proc = _fit([self.level({"method": "zero"}, values=[2])]) + (values,) = proc.process({"values": [[1.0, NAN], [NAN, 4.0]]}) + self.assertEqual(values.tolist(), [[1.0, 0.0], [0.0, 4.0]]) + + def test_statistics_by_tokens_and_dims(self): + train: Records = { + "tokens": ["A", "A", "B"], + "values": [[1.0, 10.0], [3.0, NAN], [5.0, 50.0]], + } + test: Records = {"tokens": ["A", "B", "C"], "values": [[NAN, NAN]] * 3} + cases: list[tuple[Impute, list[list[float]]]] = [ + # One mean of everything: (1 + 10 + 3 + 5 + 50) / 5. + ({"method": "mean"}, [[13.8, 13.8]] * 3), + # One mean per value position. + ({"method": "mean", "by_dims": ()}, [[3.0, 30.0]] * 3), + # One per token. "C" was never seen, so it keeps its NaN, which the + # chained "zero" step then fills. + ( + {"method": "max", "by_tokens": (0,)}, + [[10.0, 10.0], [50.0, 50.0], [0.0, 0.0]], + ), + ( + {"method": "min", "by_tokens": (0,), "by_dims": ()}, + [[1.0, 10.0], [5.0, 50.0], [0.0, 0.0]], + ), + ] + for step, expected in cases: + with self.subTest(step=step): + level = self.level(step, {"method": "zero"}, tokens=1, values=[2]) + _, values = _fit([level], train).process(test) + torch.testing.assert_close(values, torch.tensor(expected)) + + def test_statistics_come_from_raw_values(self): + # Resampling averages 1 and 5 into 3, but the max is taken from raw + # values, so the empty step 1 becomes 5. + level = self.level( + {"method": "max"}, + ordering="series", + values=[1], + resample={"step": 1.0, "method": "avg"}, + ) + raw: Records = {"time": [0.0, 0.5, 2.0], "values": [[1.0], [5.0], [2.0]]} + (values,) = _fit([level], raw).process(raw) + self.assertEqual(values.tolist(), [[3.0], [5.0], [2.0]]) + + def test_fills_never_cross_groups(self): + raw: Records = { + "children": [ + {"values": [[1.0], [NAN], [3.0]]}, + {"values": [[NAN], [2.0]]}, + ] + } + cases: list[tuple[Impute, list[float]]] = [ + ({"method": "forward_fill"}, [1.0, 1.0, 3.0, 0.0, 2.0]), + ({"method": "backward_fill"}, [1.0, 3.0, 3.0, 2.0, 2.0]), + ({"method": "linear"}, [1.0, 2.0, 3.0, 0.0, 2.0]), + ] + for step, expected in cases: + with self.subTest(step=step): + level = self.level(step, {"method": "zero"}, values=[1]) + values, _ = _fit([level, {}], raw).process(raw) + self.assertEqual(values.reshape(-1).tolist(), expected) + + def test_fills_by_token(self): + raw: Records = { + "time": [0.0, 1.0, 2.0, 4.0], + "tokens": ["A", "B", "A", "A"], + "values": [[1.0], [10.0], [NAN], [4.0]], + } + cases: list[tuple[Impute, float]] = [ + ({"method": "forward_fill", "by_tokens": (0,)}, 1.0), + ({"method": "backward_fill", "by_tokens": (0,)}, 4.0), + ({"method": "linear", "by_tokens": (0,)}, 2.5), # by time: 2 of 0..4 + ({"method": "forward_fill"}, 10.0), # any token: the previous record + ] + for step, expected in cases: + with self.subTest(step=step): + level = self.level(step, ordering="series", tokens=1, values=[1]) + _, values, _ = _fit([level], raw).process(raw) + self.assertAlmostEqual(values[2, 0].item(), expected) + + def test_fill_also_fills_pad_tokens(self): + level = self.level( + {"method": "forward_fill"}, + ordering="series", + tokens=1, + values=[1], + resample={"step": 1.0, "method": "last"}, + ) + raw: Records = { + "time": [0.0, 2.0], + "tokens": ["A", "B"], + "values": [[1.0], [2.0]], + } + tokens, values = _fit([level], raw).process(raw) + self.assertEqual(tokens.tolist(), [2, 2, 3]) # step 1 takes "A" + self.assertEqual(values.reshape(-1).tolist(), [1.0, 1.0, 2.0]) + + def test_linear_on_a_resampled_grid_uses_steps(self): + level = self.level( + {"method": "linear"}, + ordering="series", + values=[1], + resample={"step": 1.0, "method": "last"}, + ) + raw: Records = {"time": [0.0, 3.5], "values": [[0.0], [3.0]]} + (values,) = _fit([level], raw).process(raw) + self.assertEqual(values.reshape(-1).tolist(), [0.0, 1.0, 2.0, 3.0]) + + def test_unfilled_gaps_still_raise(self): + level = self.level({"method": "forward_fill"}, values=[1]) + with self.assertRaisesRegex(ValueError, "NaN"): + _fit([level]).process({"values": [[NAN], [1.0]]}) # leading gap + + def test_pickle_keeps_statistics(self): + level = self.level({"method": "mean"}, values=[1]) + proc = _fit([level], {"values": [[2.0], [4.0]]}) + restored = pickle.loads(pickle.dumps(proc)) + self.assertEqual(restored.process({"values": [[NAN]]})[0].tolist(), [[3.0]]) + + class TestCollateAndValidate(unittest.TestCase): def test_collate_nested(self): a, b = _visits(["A", "B"], ["C"]), _visits(["C"], [], ["A"]) From 2b0512b7d4183c0202231c066109fbd3939ba90f Mon Sep 17 00:00:00 2001 From: Logiquo Date: Wed, 28 Oct 2026 08:22:17 -0400 Subject: [PATCH 04/11] Add normalization to SequenceProcessorV2 - Normalize: {"method", "by_tokens", "by_dims"} with zscore or minmax, run after resampling and imputation. Statistics are learned in fit() from the raw values, per token in by_tokens and reduced over the value axes in by_dims (default (): one statistic per value position). A token without statistics (, , or only NaN in training) uses the statistics over all tokens; a constant divides by 1. - _KeyedStatistics holds per-token-key dataset statistics (count, mean, M2 via Chan et al., min, max) and is shared by _Imputer and the new _Normalizer. - normalize is no longer reserved; Normalize is exported. --- ...yhealth.processors.SequenceProcessorV2.rst | 11 +- pyhealth/processors/__init__.py | 2 + pyhealth/processors/sequence_processor_v2.py | 433 +++++++++++++----- tests/core/test_sequence_processor_v2.py | 113 ++++- 4 files changed, 424 insertions(+), 135 deletions(-) diff --git a/docs/api/processors/pyhealth.processors.SequenceProcessorV2.rst b/docs/api/processors/pyhealth.processors.SequenceProcessorV2.rst index a064c5b0f..202e1da48 100644 --- a/docs/api/processors/pyhealth.processors.SequenceProcessorV2.rst +++ b/docs/api/processors/pyhealth.processors.SequenceProcessorV2.rst @@ -35,10 +35,10 @@ In a task schema the configuration is plain data: } Batches are packed with ``collate()``, which follows ``batch_schema()``: the -fields of ``schema()`` plus the batch offsets ``[B + 1]``. ``impute`` fills NaN -values (see ``Impute``); ``normalize`` and ``missing_indicator`` are reserved -for future work and only accept ``None`` / ``False``. Values still containing -NaN or inf raise. +fields of ``schema()`` plus the batch offsets ``[B + 1]``. After resampling, +``impute`` fills NaN values (see ``Impute``) and ``normalize`` scales them (see +``Normalize``); ``missing_indicator`` is reserved for future work and only +accepts ``False``. Values still containing NaN or inf raise. .. autoclass:: pyhealth.processors.SequenceProcessorV2 :members: @@ -56,3 +56,6 @@ NaN or inf raise. .. autoclass:: pyhealth.processors.Impute :members: + +.. autoclass:: pyhealth.processors.Normalize + :members: diff --git a/pyhealth/processors/__init__.py b/pyhealth/processors/__init__.py index 28b3a15fe..507fb0ffa 100644 --- a/pyhealth/processors/__init__.py +++ b/pyhealth/processors/__init__.py @@ -40,6 +40,7 @@ def get_processor(name: str): from .sequence_processor_v2 import ( Impute, LevelConfig, + Normalize, Records, Resample, SequenceProcessorV2, @@ -87,6 +88,7 @@ def get_processor(name: str): "Records", "LevelConfig", "Impute", + "Normalize", "Resample", "SignalProcessor", "StageNetProcessor", diff --git a/pyhealth/processors/sequence_processor_v2.py b/pyhealth/processors/sequence_processor_v2.py index d421ef869..139f354eb 100644 --- a/pyhealth/processors/sequence_processor_v2.py +++ b/pyhealth/processors/sequence_processor_v2.py @@ -3,7 +3,7 @@ See ``notes/Unified_Sequence_Processor.md`` (Section 7) for the design. """ -from collections.abc import Iterable, Iterator, Sequence +from collections.abc import Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass from itertools import accumulate, repeat from typing import ( @@ -34,6 +34,7 @@ ImputeMethod = Literal[ "zero", "mean", "min", "max", "forward_fill", "backward_fill", "linear" ] +NormalizeMethod = Literal["zscore", "minmax"] # Imputation methods that use statistics learned in fit(). _STATISTICS = ("mean", "min", "max") @@ -41,7 +42,7 @@ _ALONG_TIME = ("forward_fill", "backward_fill", "linear") # Options reserved in the configuration but not implemented in Phase 1. -_RESERVED = {"normalize": None, "missing_indicator": False} +_RESERVED = {"missing_indicator": False} class Records(TypedDict, total=False): @@ -139,12 +140,46 @@ class Impute(TypedDict, total=False): by_dims: Sequence[int] | None +class Normalize(TypedDict, total=False): + """Normalization of a level's values; a plain ``dict``. + + It runs after resampling and imputation, so imputed values are scaled + too: ``mean`` imputation followed by ``zscore`` with the same keys puts + imputed values at 0. + + Keys: + method: ``"zscore"``, ``(x - mean) / std``, or ``"minmax"``, + ``(x - min) / (max - min)``. The statistics are learned in + ``fit()`` from the raw training values (before resampling), with + NaN skipped. A spread of 0 (a constant) divides by 1. + by_tokens: Token namespaces; the statistics are kept per token in + these namespaces. A token without statistics uses the statistics + over all tokens, so its value keeps its measured size on the + population scale. That is ```` (any token not seen in + ``fit()``), ```` (a resampled step without records, whose + value is NaN unless imputed), and a token whose training values + were all NaN. None (the default) ignores tokens. + by_dims: The axes of the value shape ``S`` to reduce over, indexed + like ``values`` and ``spatial``. ``()`` (the default) keeps one + statistic per value position, e.g. per channel of ``[N, F]`` + vitals; ``(1, 2)`` gives one per channel of ``[C, H, W]`` + images; None reduces over all axes. Time is always reduced. + + A value position without any training value is left as it is. + """ + + method: Required[NormalizeMethod] + by_tokens: Sequence[int] | None + by_dims: Sequence[int] | None + + class LevelConfig(TypedDict, total=False): """Configuration of one level; a plain ``dict`` with type hints. Like :class:`Records`, ``tokens`` and ``values`` declare a role by being present. They are fixed by the task, so :meth:`SequenceProcessorV2.fit` - only builds vocabularies (and impute statistics). + only builds vocabularies (and the statistics of imputation and + normalization). Keys: ordering: ``"set"``, ``"sequence"`` (default) or ``"series"``. @@ -163,7 +198,7 @@ class LevelConfig(TypedDict, total=False): spatial: One bool per value axis; defaults to all False. impute: An :class:`Impute` step, or a list of them applied in order, run after resampling. - normalize: Reserved for future work; only None. + normalize: A :class:`Normalize`, run after imputation. missing_indicator: Reserved for future work; only False. """ @@ -174,7 +209,7 @@ class LevelConfig(TypedDict, total=False): resample: Resample | None spatial: Sequence[bool] | None impute: Impute | Sequence[Impute] | None - normalize: None + normalize: Normalize | None missing_indicator: Literal[False] @@ -186,6 +221,14 @@ class _ImputeStep(NamedTuple): by_dims: tuple[int, ...] | None +class _NormalizeStep(NamedTuple): + """A checked :class:`Normalize`.""" + + method: NormalizeMethod + by_tokens: tuple[int, ...] | None + by_dims: tuple[int, ...] | None + + @dataclass(frozen=True) class _InternalLevelConfig: """A checked :class:`LevelConfig` with every default filled in. @@ -202,6 +245,7 @@ class _InternalLevelConfig: resample: tuple[float, ResampleMethod] | None # (step, method) spatial: tuple[bool, ...] | None impute: tuple[_ImputeStep, ...] + normalize: _NormalizeStep | None @staticmethod def _is_count(value: Any) -> bool: @@ -244,6 +288,42 @@ def as_pair(entry: Any) -> CodeMapping | None: f"with one pair (or None) per token namespace, got {mapping!r}." ) + @classmethod + def _parse_by( + cls, + level: int, + step: Mapping[str, object], + tokens: int, + values: tuple[int, ...], + default_dims: tuple[int, ...] | None, + ) -> tuple[tuple[int, ...] | None, tuple[int, ...] | None]: + """Checks the ``by_tokens`` and ``by_dims`` of an impute or normalize step.""" + by_tokens = step.get("by_tokens") + if by_tokens is not None: + if ( + not isinstance(by_tokens, (tuple, list)) + or not by_tokens + or not all(cls._is_index(k, tokens) for k in by_tokens) + ): + raise ValueError( + f"Level {level}: by_tokens must be None or a non-empty list of " + f"token namespaces below {tokens}, got {by_tokens!r}." + ) + by_tokens = tuple(by_tokens) + by_dims = step.get("by_dims", default_dims) + if by_dims is not None: + if ( + not isinstance(by_dims, (tuple, list)) + or not all(cls._is_index(d, len(values)) for d in by_dims) + or len(set(by_dims)) != len(by_dims) + ): + raise ValueError( + f"Level {level}: by_dims must be None or a list of distinct " + f"value axes below {len(values)}, got {by_dims!r}." + ) + by_dims = tuple(by_dims) + return by_tokens, by_dims + @classmethod def parse(cls, level: int, config: LevelConfig) -> "_InternalLevelConfig": """Checks one level's configuration and fills in the defaults. @@ -381,29 +461,7 @@ def parse(cls, level: int, config: LevelConfig) -> "_InternalLevelConfig": f"Level {level}: impute method must be one of " f"{get_args(ImputeMethod)}, got {method!r}." ) - by_tokens, by_dims = step.get("by_tokens"), step.get("by_dims") - if by_tokens is not None: - if ( - not isinstance(by_tokens, (tuple, list)) - or not by_tokens - or not all(cls._is_index(k, tokens) for k in by_tokens) - ): - raise ValueError( - f"Level {level}: by_tokens must be None or a non-empty list of " - f"token namespaces below {tokens}, got {by_tokens!r}." - ) - by_tokens = tuple(by_tokens) - if by_dims is not None: - if ( - not isinstance(by_dims, (tuple, list)) - or not all(cls._is_index(d, len(values)) for d in by_dims) - or len(set(by_dims)) != len(by_dims) - ): - raise ValueError( - f"Level {level}: by_dims must be None or a list of distinct " - f"value axes below {len(values)}, got {by_dims!r}." - ) - by_dims = tuple(by_dims) + by_tokens, by_dims = cls._parse_by(level, step, tokens, values, None) if method in (*_STATISTICS, "linear") and not values: raise ValueError( f"Level {level}: impute method {method!r} needs values." @@ -415,6 +473,27 @@ def parse(cls, level: int, config: LevelConfig) -> "_InternalLevelConfig": ) parsed.append(_ImputeStep(method, by_tokens, by_dims)) + normalize = config.get("normalize") + if normalize is not None: + if ( + not isinstance(normalize, dict) + or "method" not in normalize + or not normalize.keys() <= Normalize.__annotations__.keys() + ): + raise ValueError( + f"Level {level}: normalize must be a dict with 'method' and " + f"optionally 'by_tokens' and 'by_dims', got {normalize!r}." + ) + if normalize["method"] not in get_args(NormalizeMethod): + raise ValueError( + f"Level {level}: normalize method must be one of " + f"{get_args(NormalizeMethod)}, got {normalize['method']!r}." + ) + if not values: + raise ValueError(f"Level {level}: normalize needs values.") + by_tokens, by_dims = cls._parse_by(level, normalize, tokens, values, ()) + normalize = _NormalizeStep(normalize["method"], by_tokens, by_dims) + return cls( ordering=ordering, tokens=tokens, @@ -423,6 +502,7 @@ def parse(cls, level: int, config: LevelConfig) -> "_InternalLevelConfig": resample=resample, spatial=spatial, impute=tuple(parsed), + normalize=normalize, ) @@ -497,102 +577,189 @@ def __call__( return resampled, grid_children +def _split_by_key(keys: np.ndarray) -> Iterator[tuple[tuple[int, ...], np.ndarray]]: + """Yields each distinct row of ``keys`` and the mask of rows equal to it.""" + if len(keys) == 0: + return + if keys.shape[1] == 0: + yield (), np.ones(len(keys), dtype=bool) + return + unique, inverse = np.unique(keys, axis=0, return_inverse=True) + inverse = inverse.reshape(-1) + for index, key in enumerate(unique.tolist()): + yield tuple(key), inverse == index + + +class _KeyedStatistics: + """Dataset statistics of a level's raw values, per token key. + + The key of a record is its token IDs in the ``by_tokens`` namespaces + (``()`` when ``by_tokens`` is None), and the ``by_dims`` axes of ``S`` + are reduced. Each statistic has shape ``S`` with size 1 on the reduced + axes, so it broadcasts over the values it applies to. + + :meth:`accumulate` adds one ``Records``' values during ``fit()``, + :meth:`finish` turns the running state into statistics, and + :meth:`lookup` finds them for the records of a level in ``process()``. + With ``pooled``, the statistics over all tokens are kept as well, under + the key None, for records whose key has none. + """ + + def __init__( + self, + level: int, + tokens: int, + by_tokens: tuple[int, ...] | None, + by_dims: tuple[int, ...] | None, + pooled: bool = False, + ): + self.level, self.tokens = level, tokens + self.by_tokens, self.by_dims = by_tokens, by_dims + self.pooled = pooled and by_tokens is not None + # Running state while fitting: count, mean, M2 (Chan et al.), min, max. + self._state: dict[tuple[int, ...] | None, list[np.ndarray]] = {} + self._stats: dict[tuple[int, ...] | None, dict[str, np.ndarray]] = {} + + def _keys(self, tokens: np.ndarray | None, n: int) -> np.ndarray: + if self.by_tokens is None or tokens is None: + return np.zeros((n, 0), dtype=np.int64) + return tokens[:, list(self.by_tokens)] + + def accumulate(self, values: np.ndarray, ids: list[list[int]]) -> None: + """Adds one ``Records``' raw values ``[N, *S]`` and token IDs ``[N][K]``.""" + tokens = None + if self.by_tokens is not None: + if len(ids) != len(values): + raise ValueError( + f"Level {self.level}: {len(ids)} tokens but {len(values)} values." + ) + tokens = np.array(ids, dtype=np.int64).reshape(len(ids), self.tokens) + values = values.astype(np.float64) # sums over a dataset need the range + reduced = range(values.ndim - 1) if self.by_dims is None else self.by_dims + axes = (0, *(1 + d for d in reduced)) + groups: list[tuple[tuple[int, ...] | None, np.ndarray]] = list( + _split_by_key(self._keys(tokens, len(values))) + ) + if self.pooled and len(values): + groups.append((None, np.ones(len(values), dtype=bool))) + for key, rows in groups: + x = values[rows] + ok = np.isfinite(x) + count = ok.sum(axis=axes, keepdims=True) + with np.errstate(invalid="ignore", divide="ignore"): + mean = np.where(ok, x, 0.0).sum(axis=axes, keepdims=True) / count + mean = np.where(count > 0, mean, 0.0) + m2 = np.where(ok, (x - mean) ** 2, 0.0).sum(axis=axes, keepdims=True) + low = np.where(ok, x, np.inf).min(axis=axes, keepdims=True) + high = np.where(ok, x, -np.inf).max(axis=axes, keepdims=True) + batch = [a[0] for a in (count.astype(np.float64), mean, m2, low, high)] + if key not in self._state: + self._state[key] = batch + continue + n_a, mean_a, m2_a, low_a, high_a = self._state[key] + n_b, mean_b, m2_b, low_b, high_b = batch + total = n_a + n_b + with np.errstate(invalid="ignore", divide="ignore"): + share = np.where(total > 0, n_b / total, 0.0) + delta = mean_b - mean_a + self._state[key] = [ + total, + mean_a + delta * share, + m2_a + m2_b + delta**2 * n_a * share, + np.minimum(low_a, low_b), + np.maximum(high_a, high_b), + ] + + def finish(self) -> None: + """Turns the running state into statistics, ready for a new fit().""" + self._stats = {} + for key, (count, mean, m2, low, high) in self._state.items(): + seen = count > 0 + with np.errstate(invalid="ignore", divide="ignore"): + std = np.sqrt(m2 / count) + self._stats[key] = { + "mean": np.where(seen, mean, np.nan).astype(np.float32), + "std": np.where(seen, std, np.nan).astype(np.float32), + "min": np.where(seen, low, np.nan).astype(np.float32), + "max": np.where(seen, high, np.nan).astype(np.float32), + } + self._state = {} + + def lookup( + self, tokens: np.ndarray | None, n: int + ) -> Iterator[tuple[np.ndarray, dict[str, np.ndarray]]]: + """Yields the mask of rows that share a key, and the key's statistics. + + Rows whose key has no statistics are skipped, or, with ``pooled``, get + the statistics over all tokens. + """ + for key, rows in _split_by_key(self._keys(tokens, n)): + stats = self._stats.get(key) + if stats is None and self.pooled: + stats = self._stats.get(None) + if stats is not None: + yield rows, stats + + +class _Normalizer: + """Normalizes one level's values (see :class:`Normalize`).""" + + learns_statistics = True + + def __init__(self, level: int, step: _NormalizeStep, tokens: int): + self.method = step.method + self.statistics = _KeyedStatistics( + level, tokens, step.by_tokens, step.by_dims, pooled=True + ) + + def accumulate(self, values: np.ndarray, ids: list[list[int]]) -> None: + self.statistics.accumulate(values, ids) + + def finish(self) -> None: + self.statistics.finish() + + def __call__(self, columns: dict[Role, np.ndarray]) -> None: + """Normalizes the level's concatenated values in place.""" + values = columns["values"].astype(np.float64) + for rows, stats in self.statistics.lookup(columns.get("tokens"), len(values)): + if self.method == "zscore": + center, spread = stats["mean"], stats["std"] + else: + center, spread = stats["min"], stats["max"] - stats["min"] + # A constant divides by 1; a position never seen in fit() is kept. + spread = np.where(spread > 0, spread, 1.0) + center = np.where(np.isnan(center), 0.0, center) + values[rows] = (values[rows] - center) / spread + columns["values"] = values.astype(np.float32) + + class _Imputer: """Runs the impute steps of one level (see :class:`Impute`). - It also holds the statistics of the ``mean`` / ``min`` / ``max`` steps: - :meth:`accumulate` adds one ``Records``' raw values during ``fit()``, and - :meth:`finish` turns the running state into statistics. + It also holds the statistics of the ``mean`` / ``min`` / ``max`` steps + (:class:`_KeyedStatistics`): :meth:`accumulate` adds one ``Records``' raw + values during ``fit()``, and :meth:`finish` turns them into statistics. """ def __init__(self, level: int, steps: tuple[_ImputeStep, ...], tokens: int): - self.level, self.steps, self.tokens = level, steps, tokens - self.learns_statistics = any(step.method in _STATISTICS for step in steps) - # Per statistic step: its running state during fit(), then, per token - # key, its statistic. - self._state = self._new_state() - self._stats: dict[int, dict[tuple[int, ...], np.ndarray]] = {} - - def _new_state(self) -> dict[int, dict[tuple[int, ...], list[np.ndarray]]]: - return { - index: {} - for index, step in enumerate(self.steps) + self.level, self.steps = level, steps + # The statistics of each mean / min / max step, by step index. + self.statistics = { + index: _KeyedStatistics(level, tokens, step.by_tokens, step.by_dims) + for index, step in enumerate(steps) if step.method in _STATISTICS } + self.learns_statistics = bool(self.statistics) def accumulate(self, values: np.ndarray, ids: list[list[int]]) -> None: """Adds one ``Records``' raw values (and their token IDs) to the statistics.""" - for index, state in self._state.items(): - step = self.steps[index] - keys = np.zeros((len(values), 0), dtype=np.int64) - if step.by_tokens is not None: - if len(ids) != len(values): - raise ValueError( - f"Level {self.level}: {len(ids)} tokens but {len(values)} values." - ) - table = np.array(ids, dtype=np.int64).reshape(len(ids), self.tokens) - keys = table[:, list(step.by_tokens)] - self._accumulate(step, values, keys, state) + for statistics in self.statistics.values(): + statistics.accumulate(values, ids) def finish(self) -> None: """Turns the running state into statistics, ready for a new fit().""" - for index, state in self._state.items(): - stats: dict[tuple[int, ...], np.ndarray] = {} - for key, parts in state.items(): - if self.steps[index].method == "mean": - with np.errstate(invalid="ignore", divide="ignore"): - stat = parts[0] / parts[1] - else: - stat = np.where(np.isinf(parts[0]), np.nan, parts[0]) - stats[key] = stat.astype(np.float32) - self._stats[index] = stats - self._state = self._new_state() - - @staticmethod - def _split_by_key(keys: np.ndarray) -> Iterator[tuple[tuple[int, ...], np.ndarray]]: - """Yields each distinct row of ``keys`` and the mask of rows equal to it.""" - if len(keys) == 0: - return - if keys.shape[1] == 0: - yield (), np.ones(len(keys), dtype=bool) - return - unique, inverse = np.unique(keys, axis=0, return_inverse=True) - inverse = inverse.reshape(-1) - for index, key in enumerate(unique.tolist()): - yield tuple(key), inverse == index - - def _accumulate( - self, - step: _ImputeStep, - values: np.ndarray, - keys: np.ndarray, - state: dict[tuple[int, ...], list[np.ndarray]], - ) -> None: - """Adds raw values to the running state of one statistic, per token key.""" - reduced = range(values.ndim - 1) if step.by_dims is None else step.by_dims - axes = (0, *(1 + d for d in reduced)) - values = values.astype(np.float64) # sums over a dataset need the range - finite = np.isfinite(values) - for key, rows in self._split_by_key(keys): - x, ok = values[rows], finite[rows] - if step.method == "mean": - total = np.where(ok, x, 0.0).sum(axis=axes, keepdims=True)[0] - count = ok.sum(axis=axes, keepdims=True)[0] - if key in state: - state[key][0] += total - state[key][1] += count - else: - state[key] = [total, count] - else: - low = step.method == "min" - fill = np.inf if low else -np.inf - reduce = np.min if low else np.max - extreme = reduce(np.where(ok, x, fill), axis=axes, keepdims=True)[0] - if key in state: - combine = np.minimum if low else np.maximum - state[key][0] = combine(state[key][0], extreme) - else: - state[key] = [extreme] + for statistics in self.statistics.values(): + statistics.finish() @staticmethod def _previous_valid(valid: np.ndarray, start: np.ndarray) -> np.ndarray: @@ -653,13 +820,10 @@ def __call__(self, columns: dict[Role, np.ndarray]) -> None: ) if step.method in _STATISTICS: assert values is not None # the configuration requires values - stats = self._stats.get(index, {}) values = values.copy() - for key, rows in self._split_by_key(keys): - stat = stats.get(key) - if stat is not None: - x = values[rows] - values[rows] = np.where(np.isnan(x), stat, x) + for rows, stats in self.statistics[index].lookup(tokens, n): + x = values[rows] + values[rows] = np.where(np.isnan(x), stats[step.method], x) columns["values"] = values continue @@ -728,7 +892,7 @@ class SequenceProcessorV2(FeatureProcessor): into level ``l + 1`` records, in CSR layout. Which roles a level has, the number of token namespaces ``K`` and the value shape ``S`` come from the configuration; :meth:`fit` only builds the vocabularies and the statistics - of the impute steps. + of imputation and normalization. Subclasses are configuration only: their ``__init__`` passes fixed ``levels`` to ``super().__init__()`` and adds no logic. @@ -783,14 +947,20 @@ def __init__(self, levels: Sequence[LevelConfig]): for level, config in enumerate(self._config) ] # Fitted state: one vocabulary per level and token namespace, and the - # statistics held by each level's imputer. + # statistics held by each level's imputer and normalizer. self._imputers = [ _Imputer(level, config.impute, config.tokens) if config.impute else None for level, config in enumerate(self._config) ] + self._normalizers = [ + None + if config.normalize is None + else _Normalizer(level, config.normalize, config.tokens) + for level, config in enumerate(self._config) + ] self.learns_statistics = any( - imputer is not None and imputer.learns_statistics - for imputer in self._imputers + learner is not None and learner.learns_statistics + for learner in (*self._imputers, *self._normalizers) ) self._vocabs = [ [_new_vocab() for _ in range(config.tokens)] for config in self._config @@ -991,7 +1161,7 @@ def _empty(self, level: int, role: Role) -> np.ndarray: # ------------------------------------------------------------------ def fit(self, samples: Iterable[dict[str, Any]], field: str) -> None: - """Builds the vocabularies and the statistics of the impute steps. + """Builds the vocabularies and the statistics of imputation and normalization. Vocabularies come from the (code-mapped) tokens, statistics from the raw values. Records are not checked here; :meth:`process` checks them. @@ -1005,7 +1175,11 @@ def fit(self, samples: Iterable[dict[str, Any]], field: str) -> None: level_records: list[Records] = [sample[field]] for level in range(self.depth, -1, -1): vocabs, mappers = self._vocabs[level], self._mappers[level] - imputer = self._imputers[level] + learners = [ + learner + for learner in (self._imputers[level], self._normalizers[level]) + if learner is not None and learner.learns_statistics + ] children: list[Records] = [] for records in level_records: ids: list[list[int]] = [] @@ -1021,16 +1195,18 @@ def fit(self, samples: Iterable[dict[str, Any]], field: str) -> None: vocab[target] = len(vocab) row.append(vocab[(targets or [code])[0]]) ids.append(row) - if imputer is not None and imputer.learns_statistics: + if learners: # Statistics are taken from the raw values, before # resampling. - imputer.accumulate(self._read_values(level, records), ids) + values = self._read_values(level, records) + for learner in learners: + learner.accumulate(values, ids) children.extend(records.get("children", ())) level_records = children - for imputer in self._imputers: - if imputer is not None: - imputer.finish() + for learner in (*self._imputers, *self._normalizers): + if learner is not None: + learner.finish() self._fitted = True def process(self, value: Records) -> tuple[torch.Tensor, ...]: @@ -1118,6 +1294,9 @@ def process(self, value: Records) -> tuple[torch.Tensor, ...]: columns, children = resampler(columns, children) if imputer is not None: imputer(columns) + normalizer = self._normalizers[level] + if normalizer is not None: + normalizer(columns) for role in self._output_roles(level): array = columns[role] if role == "offsets": diff --git a/tests/core/test_sequence_processor_v2.py b/tests/core/test_sequence_processor_v2.py index e12ae5726..baeb983e0 100644 --- a/tests/core/test_sequence_processor_v2.py +++ b/tests/core/test_sequence_processor_v2.py @@ -103,10 +103,7 @@ def test_rejects_bad_configurations(self): SequenceProcessorV2(levels=levels) def test_reserved_options_not_implemented(self): - reserved: list[tuple[str, Any]] = [ - ("normalize", "zscore"), - ("missing_indicator", True), - ] + reserved: list[tuple[str, Any]] = [("missing_indicator", True)] for key, value in reserved: config: Any = {"tokens": 1, key: value} # deliberately invalid with self.subTest(key=key), self.assertRaises(NotImplementedError): @@ -772,6 +769,114 @@ def test_pickle_keeps_statistics(self): self.assertEqual(restored.process({"values": [[NAN]]})[0].tolist(), [[3.0]]) +class TestNormalize(unittest.TestCase): + @staticmethod + def level(normalize: Any, **config: Any) -> LevelConfig: + # Typed Any: the tests also build invalid configurations. + level: Any = {"normalize": normalize} + level.update(config) + return level + + def test_rejects_bad_configurations(self): + level = self.level + bad: list[Any] = [ + [level("zscore", values=[1])], + [level({"by_dims": ()}, values=[1])], # no method + [level({"method": "robust"}, values=[1])], + [level({"method": "zscore", "extra": 1}, values=[1])], + [level({"method": "zscore"}, tokens=1)], # no values + [level({"method": "zscore", "by_tokens": [1]}, tokens=1, values=[1])], + [level({"method": "zscore", "by_dims": [0, 0]}, values=[2, 2])], + ] + for levels in bad: + with self.subTest(levels=levels), self.assertRaises(ValueError): + SequenceProcessorV2(levels=levels) + + def test_learns_statistics(self): + proc = SequenceProcessorV2( + levels=[self.level({"method": "minmax"}, values=[1])] + ) + self.assertTrue(proc.learns_statistics) + + def test_zscore_per_channel_by_default(self): + raw: Records = {"values": [[1.0, 10.0], [3.0, 30.0]]} + proc = _fit([self.level({"method": "zscore"}, values=[2])], raw) + (values,) = proc.process(raw) + self.assertEqual(values.tolist(), [[-1.0, -1.0], [1.0, 1.0]]) + + def test_minmax_over_all_axes(self): + raw: Records = {"values": [[1.0, 10.0], [3.0, 21.0]]} + level = self.level({"method": "minmax", "by_dims": None}, values=[2]) + (values,) = _fit([level], raw).process(raw) + torch.testing.assert_close(values, torch.tensor([[0.0, 0.45], [0.1, 1.0]])) + + def test_statistics_stream_over_samples(self): + # Chan et al.'s update across samples equals one pass over all values. + samples: list[Records] = [ + {"values": [[1.0], [2.0]]}, + {"values": [[3.0]]}, + {"values": []}, + {"values": [[4.0], [NAN]]}, + ] + proc = _fit([self.level({"method": "zscore"}, values=[1])], *samples) + (values,) = proc.process({"values": [[2.5]]}) + self.assertAlmostEqual(values.item(), 0.0, places=6) + (values,) = proc.process({"values": [[2.5 + 1.25**0.5]]}) + self.assertAlmostEqual(values.item(), 1.0, places=6) + + def test_by_tokens_falls_back_to_all_tokens(self): + train: Records = { + "tokens": ["A", "A", "B", "B"], + "values": [[0.0], [2.0], [10.0], [30.0]], + } + level = self.level( + {"method": "minmax", "by_tokens": (0,)}, tokens=1, values=[1] + ) + proc = _fit([level], train) + _, values = proc.process( + {"tokens": ["A", "B", "C"], "values": [[1.0], [20.0], [15.0]]} + ) + # A: 0..2, B: 10..30, and unseen C: all tokens, 0..30. + self.assertEqual(values.reshape(-1).tolist(), [0.5, 0.5, 0.5]) + + def test_constant_divides_by_one(self): + raw: Records = {"values": [[5.0], [5.0]]} + proc = _fit([self.level({"method": "zscore"}, values=[1])], raw) + self.assertEqual(proc.process({"values": [[7.0]]})[0].tolist(), [[2.0]]) + + def test_position_never_seen_is_kept(self): + train: Records = {"values": [[1.0, NAN], [3.0, NAN]]} + proc = _fit([self.level({"method": "zscore"}, values=[2])], train) + (values,) = proc.process({"values": [[3.0, 7.0]]}) + self.assertEqual(values.tolist(), [[1.0, 7.0]]) + + def test_runs_after_imputation(self): + # mean imputation, then zscore with the same keys: the gap becomes 0. + level = self.level({"method": "zscore"}, values=[1], impute={"method": "mean"}) + proc = _fit([level], {"values": [[1.0], [3.0]]}) + (values,) = proc.process({"values": [[NAN], [3.0]]}) + self.assertEqual(values.tolist(), [[0.0], [1.0]]) + + def test_statistics_come_from_raw_values(self): + # Resampling averages 0 and 4 into 2; minmax still spans the raw 0..4. + level = self.level( + {"method": "minmax"}, + ordering="series", + values=[1], + resample={"step": 1.0, "method": "avg"}, + ) + raw: Records = {"time": [0.0, 0.5, 1.0], "values": [[0.0], [4.0], [1.0]]} + (values,) = _fit([level], raw).process(raw) + self.assertEqual(values.reshape(-1).tolist(), [0.5, 0.25]) + + def test_pickle_keeps_statistics(self): + proc = _fit( + [self.level({"method": "zscore"}, values=[1])], {"values": [[1.0], [3.0]]} + ) + restored = pickle.loads(pickle.dumps(proc)) + self.assertEqual(restored.process({"values": [[3.0]]})[0].tolist(), [[1.0]]) + + class TestCollateAndValidate(unittest.TestCase): def test_collate_nested(self): a, b = _visits(["A", "B"], ["C"]), _visits(["C"], [], ["A"]) From 7cbfa7b640ac074961d38010a989d15502a9d1d0 Mon Sep 17 00:00:00 2001 From: Logiquo Date: Wed, 28 Oct 2026 08:30:40 -0400 Subject: [PATCH 05/11] Add missing indicators to SequenceProcessorV2 With missing_indicator, a level emits a "mask" role right after "values": a bool tensor shaped like the values, True where a value was NaN after resampling and before imputation (resampled steps without records included). It needs values and impute. collate() and validate() handle it like any other role. Nothing in the configuration is reserved any more. --- ...yhealth.processors.SequenceProcessorV2.rst | 10 ++- pyhealth/processors/sequence_processor_v2.py | 59 +++++++++---- tests/core/test_sequence_processor_v2.py | 87 +++++++++++++++---- 3 files changed, 119 insertions(+), 37 deletions(-) diff --git a/docs/api/processors/pyhealth.processors.SequenceProcessorV2.rst b/docs/api/processors/pyhealth.processors.SequenceProcessorV2.rst index 202e1da48..371da95bc 100644 --- a/docs/api/processors/pyhealth.processors.SequenceProcessorV2.rst +++ b/docs/api/processors/pyhealth.processors.SequenceProcessorV2.rst @@ -7,8 +7,8 @@ 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``, ``time``, -``offsets``). ``offsets`` at level ``l`` partitions the level-``l`` records into +``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 @@ -37,8 +37,10 @@ In a task schema the configuration is plain data: 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``); ``missing_indicator`` is reserved for future work and only -accepts ``False``. Values still containing NaN or inf raise. +``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: diff --git a/pyhealth/processors/sequence_processor_v2.py b/pyhealth/processors/sequence_processor_v2.py index 139f354eb..0d4e71fb1 100644 --- a/pyhealth/processors/sequence_processor_v2.py +++ b/pyhealth/processors/sequence_processor_v2.py @@ -27,7 +27,7 @@ from pyhealth.medcode import CrossMap OrderType = Literal["set", "sequence", "series"] -Role = Literal["tokens", "values", "time", "offsets"] +Role = Literal["tokens", "values", "mask", "time", "offsets"] FieldKey = tuple[int, Role] CodeMapping = tuple[str, str] ResampleMethod = Literal["first", "last", "min", "max", "avg", "merge"] @@ -41,9 +41,6 @@ # Imputation methods that work along time. _ALONG_TIME = ("forward_fill", "backward_fill", "linear") -# Options reserved in the configuration but not implemented in Phase 1. -_RESERVED = {"missing_indicator": False} - class Records(TypedDict, total=False): """Raw input of one level: a collection of ``N`` records, stored as columns. @@ -199,7 +196,10 @@ class LevelConfig(TypedDict, total=False): impute: An :class:`Impute` step, or a list of them applied in order, run after resampling. normalize: A :class:`Normalize`, run after imputation. - missing_indicator: Reserved for future work; only False. + missing_indicator: If True, the level also emits ``mask``: a bool + tensor shaped like ``values``, True where a value was missing + (NaN, including a resampled step without records) before + imputation. Needs values and ``impute``. """ ordering: OrderType @@ -210,7 +210,7 @@ class LevelConfig(TypedDict, total=False): spatial: Sequence[bool] | None impute: Impute | Sequence[Impute] | None normalize: Normalize | None - missing_indicator: Literal[False] + missing_indicator: bool class _ImputeStep(NamedTuple): @@ -246,6 +246,7 @@ class _InternalLevelConfig: spatial: tuple[bool, ...] | None impute: tuple[_ImputeStep, ...] normalize: _NormalizeStep | None + missing_indicator: bool @staticmethod def _is_count(value: Any) -> bool: @@ -341,13 +342,6 @@ def parse(cls, level: int, config: LevelConfig) -> "_InternalLevelConfig": f"Level {level}: unknown configuration keys {sorted(unknown)}." ) - for key, default in _RESERVED.items(): - if config.get(key, default) != default: - raise NotImplementedError( - f"Level {level}: '{key}' is reserved for future work and only " - f"accepts {default!r} for now." - ) - ordering = config.get("ordering", "sequence") if ordering not in get_args(OrderType): raise ValueError( @@ -473,6 +467,21 @@ def parse(cls, level: int, config: LevelConfig) -> "_InternalLevelConfig": ) parsed.append(_ImputeStep(method, by_tokens, by_dims)) + missing_indicator = config.get("missing_indicator", False) + if not isinstance(missing_indicator, bool): + raise TypeError( + f"Level {level}: missing_indicator must be a bool, got " + f"{missing_indicator!r}." + ) + if missing_indicator and not values: + raise ValueError(f"Level {level}: missing_indicator needs values.") + if missing_indicator and not parsed: + # Without imputation any NaN raises, so the mask would be all False. + raise ValueError( + f"Level {level}: missing_indicator marks the values imputation " + "fills, so it needs impute." + ) + normalize = config.get("normalize") if normalize is not None: if ( @@ -503,6 +512,7 @@ def parse(cls, level: int, config: LevelConfig) -> "_InternalLevelConfig": spatial=spatial, impute=tuple(parsed), normalize=normalize, + missing_indicator=missing_indicator, ) @@ -887,7 +897,7 @@ class SequenceProcessorV2(FeatureProcessor): """Processor for hierarchical packed records, without padding or truncation. Reads a :class:`Records` and emits a tuple of tensors in :meth:`schema` - order: by level, then by role (``tokens``, ``values``, ``time``, + order: by level, then by role (``tokens``, ``values``, ``mask``, ``time``, ``offsets``). ``offsets`` at level ``l`` partitions level ``l`` records into level ``l + 1`` records, in CSR layout. Which roles a level has, the number of token namespaces ``K`` and the value shape ``S`` come from the @@ -998,6 +1008,7 @@ def _output_roles(self, level: int) -> tuple[Role, ...]: present = { "tokens": config.tokens > 0, "values": len(config.values) > 0, + "mask": config.missing_indicator, "time": config.ordering == "series" and config.resample is None, "offsets": level < self.depth, } @@ -1217,8 +1228,9 @@ def process(self, value: Records) -> tuple[torch.Tensor, ...]: Returns: Tuple of tensors: ``int64`` tokens (``[N]``, or ``[N, K]``), - ``float32`` values ``[N, *S]``, ``float32`` time ``[N]`` and - ``int64`` offsets ``[N_{l+1} + 1]``. + ``float32`` values ``[N, *S]``, ``bool`` mask ``[N, *S]`` (True + where a value was missing), ``float32`` time ``[N]`` and ``int64`` + offsets ``[N_{l+1} + 1]``. """ self._require_fit() unk = TokenProcessorInterface.UNK @@ -1292,6 +1304,9 @@ def process(self, value: Records) -> tuple[torch.Tensor, ...]: resampler, imputer = self._resamplers[level], self._imputers[level] if resampler is not None: columns, children = resampler(columns, children) + if config.missing_indicator: + # Missing before imputation fills it, including resampled gaps. + columns["mask"] = np.isnan(columns["values"]) if imputer is not None: imputer(columns) normalizer = self._normalizers[level] @@ -1368,7 +1383,11 @@ def validate(self, output: tuple[torch.Tensor, ...], batched: bool = False) -> N raise TypeError( f"{(level, role)}: expected a tensor, got {type(tensor).__name__}." ) - dtype = torch.long if role in ("tokens", "offsets") else torch.float32 + dtype = { + "tokens": torch.long, + "offsets": torch.long, + "mask": torch.bool, + }.get(role, torch.float32) if tensor.dtype != dtype: raise ValueError( f"{(level, role)}: expected {dtype}, got {tensor.dtype}." @@ -1391,6 +1410,12 @@ def validate(self, output: tuple[torch.Tensor, ...], batched: bool = False) -> N raise ValueError( f"{(level, role)}: token index outside the vocabulary." ) + elif role == "mask": + shape = self._config[level].values + if tuple(tensor.shape[1:]) != shape or tensor.dim() != 1 + len(shape): + raise ValueError( + f"{(level, role)}: expected [N, *{shape}], got {tuple(tensor.shape)}." + ) elif role == "values": shape = self._config[level].values if tuple(tensor.shape[1:]) != shape or tensor.dim() != 1 + len(shape): diff --git a/tests/core/test_sequence_processor_v2.py b/tests/core/test_sequence_processor_v2.py index baeb983e0..476ccf9fc 100644 --- a/tests/core/test_sequence_processor_v2.py +++ b/tests/core/test_sequence_processor_v2.py @@ -102,22 +102,6 @@ def test_rejects_bad_configurations(self): ): SequenceProcessorV2(levels=levels) - def test_reserved_options_not_implemented(self): - reserved: list[tuple[str, Any]] = [("missing_indicator", True)] - for key, value in reserved: - config: Any = {"tokens": 1, key: value} # deliberately invalid - with self.subTest(key=key), self.assertRaises(NotImplementedError): - SequenceProcessorV2(levels=[config]) - SequenceProcessorV2( - levels=[ - { - "tokens": 1, - "normalize": None, - "missing_indicator": False, - } - ] - ) - def test_only_vocabulary_needs_fit(self): proc = SequenceProcessorV2(levels=[{"tokens": 2, "values": [3]}]) self.assertEqual(proc.schema(), ((0, "tokens"), (0, "values"))) @@ -877,6 +861,77 @@ def test_pickle_keeps_statistics(self): self.assertEqual(restored.process({"values": [[3.0]]})[0].tolist(), [[1.0]]) +class TestMissingIndicator(unittest.TestCase): + def test_rejects_bad_configurations(self): + bad: list[Any] = [ + [{"values": [1], "missing_indicator": 1}], + [{"tokens": 1, "missing_indicator": True}], # no values + [{"values": [1], "missing_indicator": True}], # no impute + ] + for levels in bad: + with ( + self.subTest(levels=levels), + self.assertRaises((ValueError, TypeError)), + ): + SequenceProcessorV2(levels=levels) + + def test_mask_follows_values_in_the_schema(self): + level: LevelConfig = { + "ordering": "series", + "tokens": 1, + "values": [2], + "impute": {"method": "zero"}, + "missing_indicator": True, + } + proc = SequenceProcessorV2(levels=[level]) + self.assertEqual( + proc.schema(), ((0, "tokens"), (0, "values"), (0, "mask"), (0, "time")) + ) + + def test_mask_marks_values_missing_before_imputation(self): + level: LevelConfig = { + "values": [2], + "impute": {"method": "zero"}, + "missing_indicator": True, + } + values, mask = _fit([level]).process({"values": [[1.0, NAN], [NAN, 4.0]]}) + self.assertEqual(values.tolist(), [[1.0, 0.0], [0.0, 4.0]]) + self.assertEqual(mask.dtype, torch.bool) + self.assertEqual(mask.tolist(), [[False, True], [True, False]]) + + def test_resampled_gaps_are_missing(self): + level: LevelConfig = { + "ordering": "series", + "values": [1], + "resample": {"step": 1.0, "method": "last"}, + "impute": {"method": "forward_fill"}, + "missing_indicator": True, + } + raw: Records = {"time": [0.0, 2.0], "values": [[1.0], [3.0]]} + values, mask = _fit([level]).process(raw) + self.assertEqual(values.reshape(-1).tolist(), [1.0, 1.0, 3.0]) + self.assertEqual(mask.reshape(-1).tolist(), [False, True, False]) + + def test_collate_and_validate(self): + level: LevelConfig = { + "values": [1], + "impute": {"method": "zero"}, + "missing_indicator": True, + } + proc = _fit([level, {}]) + a = proc.process({"children": [{"values": [[NAN]]}, {"values": [[2.0]]}]}) + b = proc.process({"children": [{"values": []}]}) + proc.validate(a) + batch = proc.collate([a, b]) + proc.validate(batch, batched=True) + self.assertEqual(batch[1].reshape(-1).tolist(), [True, False]) + values, mask, offsets = a + with self.assertRaises(ValueError): + proc.validate((values, mask.float(), offsets)) # mask must be bool + with self.assertRaises(ValueError): + proc.validate((values, mask[:, :0], offsets)) # and shaped like values + + class TestCollateAndValidate(unittest.TestCase): def test_collate_nested(self): a, b = _visits(["A", "B"], ["C"]), _visits(["C"], [], ["A"]) From f8938169c585bb588b2a7ca4f44c2d87f3e7612f Mon Sep 17 00:00:00 2001 From: Logiquo Date: Wed, 28 Oct 2026 08:47:20 -0400 Subject: [PATCH 06/11] Store SequenceProcessorV2 fields as tensors and batch them with collate() - stores_tensor now also covers a tuple of tensors with the same length for every sample; litdata stores such tuples natively, so SampleBuilder.encode writes them as tensors instead of pickling them. SequenceProcessorV2 declares stores_tensor = True. - A processor can define collate(items). collate_fn_dict_with_padding takes an optional mapping of per-field collators, checked before the (time, values) heuristic, and get_dataloader passes each processor's collate for its field. Other fields and loaders are unchanged. --- pyhealth/datasets/sample_dataset.py | 13 ++++--- pyhealth/datasets/utils.py | 26 ++++++++++++-- pyhealth/processors/base_processor.py | 15 +++++--- pyhealth/processors/sequence_processor_v2.py | 4 +++ tests/core/test_sequence_processor_v2.py | 37 ++++++++++++++++++++ 5 files changed, 83 insertions(+), 12 deletions(-) diff --git a/pyhealth/datasets/sample_dataset.py b/pyhealth/datasets/sample_dataset.py index efcc47f32..5bd9d1f66 100644 --- a/pyhealth/datasets/sample_dataset.py +++ b/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: diff --git a/pyhealth/datasets/utils.py b/pyhealth/datasets/utils.py index fa0acf5f7..fce7e9034 100644 --- a/pyhealth/datasets/utils.py +++ b/pyhealth/datasets/utils.py @@ -1,6 +1,8 @@ +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 @@ -250,11 +252,15 @@ 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: Optional[Mapping[str, Callable]] = 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. @@ -267,8 +273,10 @@ def collate_fn_dict_with_padding(batch: List[dict]) -> dict: 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 +411,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) diff --git a/pyhealth/processors/base_processor.py b/pyhealth/processors/base_processor.py index 28915c89e..4be5fbc2a 100644 --- a/pyhealth/processors/base_processor.py +++ b/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 diff --git a/pyhealth/processors/sequence_processor_v2.py b/pyhealth/processors/sequence_processor_v2.py index 0d4e71fb1..6a1219b70 100644 --- a/pyhealth/processors/sequence_processor_v2.py +++ b/pyhealth/processors/sequence_processor_v2.py @@ -920,6 +920,10 @@ class SequenceProcessorV2(FeatureProcessor): (tensor([2, 3, 4]), tensor([0, 2, 2, 3])) """ + # process() returns the same number of tensors for every sample, which + # the disk cache stores as tensors. + stores_tensor = True + def __init__(self, levels: Sequence[LevelConfig]): if not isinstance(levels, (list, tuple)) or not levels: raise ValueError("levels must be a non-empty list of dicts.") diff --git a/tests/core/test_sequence_processor_v2.py b/tests/core/test_sequence_processor_v2.py index 476ccf9fc..74f41d176 100644 --- a/tests/core/test_sequence_processor_v2.py +++ b/tests/core/test_sequence_processor_v2.py @@ -1001,6 +1001,43 @@ def test_validate_rejects_broken_outputs(self): class TestStorage(unittest.TestCase): + def test_dataloader_uses_the_processor_collate(self): + from pyhealth.datasets import create_sample_dataset, get_dataloader + + deep: list[LevelConfig] = [{"ordering": "set", "tokens": 1}, {}, {}] + samples = [ + { + "patient_id": "p0", + "visits": _visits(["A", "B"], ["C"]), + "stays": {"children": [_visits(["A"])]}, + "label": 0, + }, + { + "patient_id": "p1", + "visits": _visits(["C"]), + "stays": {"children": [_visits(["B"], ["C"])]}, + "label": 1, + }, + ] + dataset = create_sample_dataset( + samples=samples, + input_schema={ + "visits": ("sequence_v2", {"levels": NESTED}), # two tensors + "stays": ("sequence_v2", {"levels": deep}), # three tensors + }, + output_schema={"label": "binary"}, + ) + batch = next(iter(get_dataloader(dataset, batch_size=2))) + for key in ("visits", "stays"): + proc = dataset.input_processors[key] + assert isinstance(proc, SequenceProcessorV2) + expected = proc.collate([dataset[0][key], dataset[1][key]]) + self.assertEqual(len(batch[key]), len(expected)) + for got, want in zip(batch[key], expected): + torch.testing.assert_close(got, want) + proc.validate(batch[key], batched=True) + self.assertEqual(batch["label"].shape[0], 2) # other fields as before + def test_pickle_round_trip(self): raw = _visits(["A", "B"], ["C"]) proc = pickle.loads(pickle.dumps(_fit(NESTED, raw))) From 90fdd9ac3758e3db53d57e1e09b2214998d9f4d7 Mon Sep 17 00:00:00 2001 From: Logiquo Date: Wed, 28 Oct 2026 08:53:29 -0400 Subject: [PATCH 07/11] Fix time-based imputation on a level without records forward_fill, backward_fill and linear reshaped the level's values with reshape(n, -1), which numpy cannot do when n is 0 (e.g. a sample whose stays have no lab records). Use the explicit value width instead, and test every impute method on such a level. --- pyhealth/processors/sequence_processor_v2.py | 3 ++- tests/core/test_sequence_processor_v2.py | 24 ++++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/pyhealth/processors/sequence_processor_v2.py b/pyhealth/processors/sequence_processor_v2.py index 6a1219b70..84a4e723d 100644 --- a/pyhealth/processors/sequence_processor_v2.py +++ b/pyhealth/processors/sequence_processor_v2.py @@ -841,7 +841,8 @@ def __call__(self, columns: dict[Role, np.ndarray]) -> None: order, start, end = self._segments(group, keys) skip = (keys[order] == pad).any(axis=1) # rows without a token if values is not None: - x = values.reshape(n, -1)[order] + # An explicit width: numpy cannot infer -1 when n is 0. + x = values.reshape(n, int(np.prod(values.shape[1:])))[order] valid = ~np.isnan(x) before = self._previous_valid(valid, start) after = self._next_valid(valid, start, end) diff --git a/tests/core/test_sequence_processor_v2.py b/tests/core/test_sequence_processor_v2.py index 74f41d176..d926b0b9b 100644 --- a/tests/core/test_sequence_processor_v2.py +++ b/tests/core/test_sequence_processor_v2.py @@ -741,6 +741,30 @@ def test_linear_on_a_resampled_grid_uses_steps(self): (values,) = _fit([level], raw).process(raw) self.assertEqual(values.reshape(-1).tolist(), [0.0, 1.0, 2.0, 3.0]) + def test_levels_without_records(self): + # A sample whose stays have no records at all, for every method. + for method in ("zero", "mean", "forward_fill", "backward_fill", "linear"): + with self.subTest(method=method): + level = self.level( + {"method": method}, ordering="series", tokens=1, values=[2] + ) + proc = _fit( + [level, {}], + { + "children": [ + {"time": [0.0], "tokens": ["A"], "values": [[1.0, 2.0]]} + ] + }, + ) + empty: Records = {"time": [], "tokens": [], "values": []} + tokens, values, time, offsets = proc.process( + {"children": [empty, empty]} + ) + self.assertEqual( + (tokens.shape, values.shape, time.shape), ((0,), (0, 2), (0,)) + ) + self.assertEqual(offsets.tolist(), [0, 0, 0]) + def test_unfilled_gaps_still_raise(self): level = self.level({"method": "forward_fill"}, values=[1]) with self.assertRaisesRegex(ValueError, "NaN"): From 51f5c7e79e959edde71640fc43343679a2e8ae12 Mon Sep 17 00:00:00 2001 From: Logiquo Date: Wed, 28 Oct 2026 09:01:59 -0400 Subject: [PATCH 08/11] Add docstring examples and modern annotations for the PR rules - Records, Resample, Impute, Normalize and LevelConfig get runnable '>>>' examples (all module doctests pass). - collate_fn_dict_with_padding gets an example and list / X | None annotations; its import block is sorted. --- pyhealth/datasets/utils.py | 15 ++++- pyhealth/processors/sequence_processor_v2.py | 69 ++++++++++++++++++++ 2 files changed, 82 insertions(+), 2 deletions(-) diff --git a/pyhealth/datasets/utils.py b/pyhealth/datasets/utils.py index fce7e9034..55da18209 100644 --- a/pyhealth/datasets/utils.py +++ b/pyhealth/datasets/utils.py @@ -6,8 +6,8 @@ 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 @@ -253,7 +253,7 @@ def collate_fn_dict(batch: List[dict]) -> dict: def collate_fn_dict_with_padding( - batch: List[dict], collators: Optional[Mapping[str, Callable]] = None + batch: list[dict], collators: Mapping[str, Callable] | None = None ) -> dict: """Collates a batch of data into a dictionary with padding for tensor values. @@ -266,6 +266,17 @@ def collate_fn_dict_with_padding( 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() diff --git a/pyhealth/processors/sequence_processor_v2.py b/pyhealth/processors/sequence_processor_v2.py index 84a4e723d..737333796 100644 --- a/pyhealth/processors/sequence_processor_v2.py +++ b/pyhealth/processors/sequence_processor_v2.py @@ -60,6 +60,20 @@ class Records(TypedDict, total=False): list is also accepted when ``S`` is ``(1,)``. children: ``[N]`` ``Records``; ``children[i]`` holds record ``i``'s records one level down. Read on every level but 0. + + Examples: + Two timed visits, each a set of codes: + + >>> visits: Records = { + ... "time": [0.0, 30.0], + ... "children": [{"tokens": ["A", "B"]}, {"tokens": ["C"]}], + ... } + >>> proc = SequenceProcessorV2( + ... levels=[{"ordering": "set", "tokens": 1}, {"ordering": "series"}] + ... ) + >>> proc.fit([{"visits": visits}], "visits") + >>> proc.process(visits) + (tensor([2, 3, 4]), tensor([0, 2, 3]), tensor([ 0., 30.])) """ time: Sequence[float] | np.ndarray @@ -90,6 +104,18 @@ class Resample(TypedDict): A step without records gets ```` tokens, NaN values and no children; filling it is up to imputation. + + Examples: + Averaging irregular values per time unit: + + >>> resample: Resample = {"step": 1.0, "method": "avg"} + >>> proc = SequenceProcessorV2( + ... levels=[{"ordering": "series", "values": [1], "resample": resample}] + ... ) + >>> proc.fit([], "x") + >>> proc.process({"time": [0.0, 0.5, 1.2], "values": [[1.0], [3.0], [5.0]]}) + (tensor([[2.], + [5.]]),) """ step: float @@ -130,6 +156,21 @@ class Impute(TypedDict, total=False): A key without a statistic (a token or value never seen in ``fit()``) keeps its NaN, as does any gap a method cannot fill; chain several steps to fill those. + + Examples: + Forward fill, then zero for the leading gap it cannot fill: + + >>> steps: list[Impute] = [{"method": "forward_fill"}, {"method": "zero"}] + >>> proc = SequenceProcessorV2( + ... levels=[{"ordering": "series", "values": [1], "impute": steps}] + ... ) + >>> proc.fit([], "x") + >>> nan = float("nan") + >>> values, time = proc.process( + ... {"time": [0.0, 1.0, 2.0], "values": [[nan], [2.0], [nan]]} + ... ) + >>> values.reshape(-1) + tensor([0., 2., 2.]) """ method: Required[ImputeMethod] @@ -163,6 +204,17 @@ class Normalize(TypedDict, total=False): images; None reduces over all axes. Time is always reduced. A value position without any training value is left as it is. + + Examples: + One z-score per channel, learned from the training values: + + >>> normalize: Normalize = {"method": "zscore"} + >>> proc = SequenceProcessorV2( + ... levels=[{"values": [2], "normalize": normalize}] + ... ) + >>> proc.fit([{"x": {"values": [[1.0, 10.0], [3.0, 30.0]]}}], "x") + >>> proc.process({"values": [[1.0, 30.0]]}) + (tensor([[-1., 1.]]),) """ method: Required[NormalizeMethod] @@ -200,6 +252,23 @@ class LevelConfig(TypedDict, total=False): tensor shaped like ``values``, True where a value was missing (NaN, including a resampled step without records) before imputation. Needs values and ``impute``. + + Examples: + Long-format labs (test, unit) in stays, with missing values imputed + and marked: + + >>> levels: list[LevelConfig] = [ + ... { + ... "ordering": "series", + ... "tokens": 2, + ... "values": [1], + ... "impute": {"method": "zero"}, + ... "missing_indicator": True, + ... }, + ... {"ordering": "sequence"}, + ... ] + >>> SequenceProcessorV2(levels=levels).schema() + ((0, 'tokens'), (0, 'values'), (0, 'mask'), (0, 'time'), (0, 'offsets')) """ ordering: OrderType From d5420d82ec48914461a27aba103b01b5b81b0854 Mon Sep 17 00:00:00 2001 From: Logiquo Date: Fri, 30 Oct 2026 10:43:29 +0000 Subject: [PATCH 09/11] Address review: reject multi-target mappings early and test new lines - A one-to-many code_mapping is now rejected in __init__ unless the level is a bottom-level 'set' with one token namespace and no values. Before, a set level with children or several namespaces was accepted and then failed in process() on the first sample with such a code; that runtime check can no longer trigger and is removed. - SampleBuilder.encode checks that a stored tuple of tensors has one tensor per schema() entry, so every sample (in every worker) has the same structure for litdata. - SequenceProcessorV2._count drops an unreachable branch, and the module docstring no longer points to a design note outside the repo. - Tests cover the new rejections, collate_fn_dict_with_padding with collators, the tuple-length check, and the remaining validate, vocab, time and fit error paths. --- pyhealth/datasets/sample_dataset.py | 22 +++++--- pyhealth/processors/base_processor.py | 10 ++-- pyhealth/processors/sequence_processor_v2.py | 57 +++++++++----------- tests/core/test_collate.py | 29 ++++++++++ tests/core/test_raw_storage.py | 29 ++++++++++ tests/core/test_sequence_processor_v2.py | 57 +++++++++++++------- 6 files changed, 144 insertions(+), 60 deletions(-) create mode 100644 tests/core/test_collate.py diff --git a/pyhealth/datasets/sample_dataset.py b/pyhealth/datasets/sample_dataset.py index 5bd9d1f66..b60702039 100644 --- a/pyhealth/datasets/sample_dataset.py +++ b/pyhealth/datasets/sample_dataset.py @@ -354,10 +354,10 @@ def encode(self, transformed: dict[str, Any]) -> dict[str, Any]: 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 (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 + tuples of tensors, one per ``schema()`` entry); 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 on read, so callers get back the original Python objects. Args: @@ -368,9 +368,11 @@ def encode(self, transformed: dict[str, Any]) -> dict[str, Any]: Raises: ValueError: If a field declared ``stores_tensor`` is not a tensor - or a tuple of tensors. + or a tuple of tensors, one per entry of its processor's + ``schema()``. """ tensor_fields = self._tensor_fields() + processors = {**self._input_processors, **self._output_processors} record: dict[str, Any] = {} objects: dict[str, Any] = {} for key, value in transformed.items(): @@ -378,14 +380,22 @@ def encode(self, transformed: dict[str, Any]) -> dict[str, Any]: is_tensors = isinstance(value, tuple) and all( isinstance(item, torch.Tensor) for item in value ) + processor = processors[key] 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 or a tuple of them." ) + # litdata expects the same structure in every sample, and + # schema() fixes the tuple's length for all of them. + if is_tensors and len(value) != len(processor.schema()): + raise ValueError( + f"Field '{key}': {type(processor).__name__}.process() " + f"returned {len(value)} tensors, but its schema() has " + f"{len(processor.schema())} entries." + ) record[key] = value else: objects[key] = value diff --git a/pyhealth/processors/base_processor.py b/pyhealth/processors/base_processor.py index 4be5fbc2a..d3449869c 100644 --- a/pyhealth/processors/base_processor.py +++ b/pyhealth/processors/base_processor.py @@ -71,11 +71,11 @@ class FeatureProcessor(Processor): (``set_task(task, split=PatientSplit(...))``). 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. + ``torch.Tensor``, or a tuple with one tensor per :meth:`schema` entry + (each tensor may change shape between samples). 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; diff --git a/pyhealth/processors/sequence_processor_v2.py b/pyhealth/processors/sequence_processor_v2.py index 737333796..e5d0d55c6 100644 --- a/pyhealth/processors/sequence_processor_v2.py +++ b/pyhealth/processors/sequence_processor_v2.py @@ -1,7 +1,4 @@ -"""Unified processor for hierarchical packed records (CSR / jagged layout). - -See ``notes/Unified_Sequence_Processor.md`` (Section 7) for the design. -""" +"""Unified processor for hierarchical packed records (CSR / jagged layout).""" from collections.abc import Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass @@ -241,7 +238,8 @@ class LevelConfig(TypedDict, total=False): ``pyhealth.medcode.CrossMap``, or a list with one pair (or None) per token namespace. A code that maps to several targets becomes several records, so a mapping with several targets per code is - only allowed on a ``set`` level without values. + only allowed on a bottom-level ``set`` with one token namespace + and no values. resample: A :class:`Resample` for a ``series`` level: puts its records on a time grid, turning it into a ``sequence`` level. spatial: One bool per value axis; defaults to all False. @@ -1007,15 +1005,26 @@ def __init__(self, levels: Sequence[LevelConfig]): ] for level, (config, mappers) in enumerate(zip(self._config, self._mappers)): # A code with several targets becomes several records, which only - # makes sense in an unordered level with no values to copy. + # makes sense in an unordered bottom level with nothing else to + # copy: no values, no children and no other namespaces. many = [m for m in mappers or () if m is not None and m.has_many_targets] - if many and (config.ordering != "set" or config.values): + problems = [ + problem + for problem, present in ( + (f"ordering {config.ordering!r}", config.ordering != "set"), + ("values", bool(config.values)), + ("children", level > 0), + (f"{config.tokens} token namespaces", config.tokens > 1), + ) + if present + ] + if many and problems: raise ValueError( f"Level {level}: {many[0].s_vocab}->{many[0].t_vocab} maps some " - "codes to several targets, which is only allowed on a 'set' " - f"level without values (got ordering {config.ordering!r}" - f"{' with values' if config.values else ''}). Put the tokens " - "to map in their own 'set' level, one level below this one." + "codes to several targets, which is only allowed on a " + "bottom-level 'set' with one token namespace and no values " + f"(got {', '.join(problems)}). Put the tokens to map in their " + "own 'set' level at the bottom." ) for level, config in enumerate(self._config): @@ -1185,19 +1194,8 @@ def _map_tokens( targets.append(mapped if mapped else [code]) if all(len(t) == 1 for t in targets): mapped_tokens.append(tuple(t[0] for t in targets)) - continue - # The configuration rules out time and values on a mapped level. - if level > 0 or len(token) > 1: - reason = ( - "its records also have children" - if level > 0 - else "there are several namespaces" - ) - raise ValueError( - f"Level {level}: {token!r} maps to several codes, which would duplicate " - f"records, but {reason}." - ) - mapped_tokens.extend((target,) for target in targets[0]) + else: # the configuration allows this only with one namespace + mapped_tokens.extend((target,) for target in targets[0]) return mapped_tokens def _read_values(self, level: int, records: Records) -> np.ndarray: @@ -1431,13 +1429,10 @@ def collate( def _count( schema: tuple[FieldKey, ...], output: tuple[torch.Tensor, ...], level: int ) -> int: - """Number of records at ``level``, read from any field that knows it.""" - for (field_level, role), tensor in zip(schema, output): - if field_level == level and role != "offsets": - return tensor.shape[0] - if field_level == level - 1 and role == "offsets": - return tensor.shape[0] - 1 - raise ValueError(f"No field records the size of level {level}.") + """Number of records at ``level``, read from the level below's offsets.""" + if level > 0: + return output[schema.index((level - 1, "offsets"))].shape[0] - 1 + return output[0].shape[0] # level 0 always has tokens, values or time def validate(self, output: tuple[torch.Tensor, ...], batched: bool = False) -> None: """Checks an output of :meth:`process` (or :meth:`collate` if ``batched``). diff --git a/tests/core/test_collate.py b/tests/core/test_collate.py new file mode 100644 index 000000000..dd4db2a0f --- /dev/null +++ b/tests/core/test_collate.py @@ -0,0 +1,29 @@ +import unittest + +import torch + +from pyhealth.datasets.utils import collate_fn_dict_with_padding + + +class TestCollateFnDictWithPadding(unittest.TestCase): + def test_collators_take_only_their_keys(self): + batch = [ + {"x": (torch.tensor([1.0]), torch.tensor([0, 1])), "y": torch.tensor([1, 2]), "id": "a"}, + {"x": (torch.tensor([2.0, 3.0]), torch.tensor([0, 2])), "y": torch.tensor([3]), "id": "b"}, + ] + calls = [] + + def collate_x(items): + calls.append(items) + return "collated" + + collated = collate_fn_dict_with_padding(batch, collators={"x": collate_x}) + self.assertEqual(collated["x"], "collated") + self.assertEqual(calls, [[batch[0]["x"], batch[1]["x"]]]) + # Keys without a collator are padded or listed as before. + self.assertEqual(collated["y"].tolist(), [[1, 2], [3, 0]]) + self.assertEqual(collated["id"], ["a", "b"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/core/test_raw_storage.py b/tests/core/test_raw_storage.py index 96368b2b3..24ce0e0f6 100644 --- a/tests/core/test_raw_storage.py +++ b/tests/core/test_raw_storage.py @@ -95,6 +95,21 @@ def process(self, value): return list(value) +class TupleProcessor(FeatureProcessor): + """Declares tensor output and returns a tuple, as `value` says.""" + + stores_tensor = True + + def __init__(self, value): + self.value = value + + def process(self, value): + return self.value + + def schema(self): + return ("value", "mask") + + def _samples(values): return [ { @@ -166,6 +181,20 @@ def test_misdeclared_tensor_processor_raises_with_field_name(self): in_memory=False, ) + def test_tensor_tuples_must_match_the_schema(self): + cases = { + "'codes'.*stores_tensor=True": (torch.zeros(1), "mask"), # not all tensors + r"'codes'.*1 tensors.*schema\(\) has 2": (torch.zeros(1),), + } + for message, value in cases.items(): + with self.subTest(message), self.assertRaisesRegex(ValueError, message): + create_sample_dataset( + _samples(CASES["scalar str"]), + {"raw_field": "raw", "codes": TupleProcessor(value)}, + OUTPUT, + in_memory=False, + ) + def test_disk_write_starts_no_processes(self): # litdata.optimize left helper processes running when the transform # raised; they deadlocked a later resource-tracker shutdown on Linux. diff --git a/tests/core/test_sequence_processor_v2.py b/tests/core/test_sequence_processor_v2.py index d926b0b9b..111d1fb31 100644 --- a/tests/core/test_sequence_processor_v2.py +++ b/tests/core/test_sequence_processor_v2.py @@ -50,6 +50,12 @@ def test_levels_are_required(self): def test_registered(self): self.assertIs(get_processor("sequence_v2"), SequenceProcessorV2) + def test_repr_shows_levels(self): + self.assertEqual( + repr(SequenceProcessorV2(levels=[CODES])), + "SequenceProcessorV2(levels=[{'tokens': 1}])", + ) + def test_subclass_is_configuration_only(self): class NestedCodes(SequenceProcessorV2): def __init__(self): @@ -111,6 +117,8 @@ def test_only_vocabulary_needs_fit(self): proc.vocab_size() with self.assertRaises(RuntimeError): proc.vocab(0) + with self.assertRaises(ValueError): + _fit([SCALARS]).vocab(0) # no tokens with self.assertRaises(RuntimeError): proc.process({"tokens": [("A", "u")], "values": [[1.0, 2.0, 3.0]]}) @@ -331,6 +339,7 @@ def test_bad_records_raise_when_read(self): ([series], {"tokens": ["A"]}), # series without time ([series], {"time": [1.0, 0.0], "tokens": ["A", "B"]}), ([series], {"time": [NAN], "tokens": ["A"]}), + ([series], {"time": [[0.0]], "tokens": ["A"]}), ([SCALARS], {"values": [[1.0, 2.0]]}), ([SCALARS], {"values": [[[1.0]]]}), ([SCALARS], {"values": [[NAN]]}), @@ -393,6 +402,7 @@ def test_one_to_many_expands_records(self): self.assertEqual(offsets.tolist(), [0, 3, 5]) def test_per_namespace_mapping(self): + self.use_table({"250.00": ["49"]}) # several namespaces need one target proc = _fit( [{"ordering": "set", "tokens": 2, "code_mapping": [None, ("A", "B")]}], {"tokens": [("401.9", "250.00")]}, @@ -400,12 +410,14 @@ def test_per_namespace_mapping(self): self.assertEqual(proc.vocab(0, 0), {"": 0, "": 1, "401.9": 2}) self.assertEqual(proc.vocab(0, 1), {"": 0, "": 1, "49": 2}) - def test_multi_target_mapping_needs_a_set_level_without_values(self): + def test_multi_target_mapping_needs_a_bottom_set_level_of_codes(self): mapping = self.MAPPING cases: list[list[LevelConfig]] = [ [{"tokens": 1, "code_mapping": mapping}], # sequence [{"ordering": "series", "tokens": 1, "code_mapping": mapping}], [{"ordering": "set", "tokens": 1, "values": [1], "code_mapping": mapping}], + [CODES, {"ordering": "set", "tokens": 1, "code_mapping": mapping}], + [{"ordering": "set", "tokens": 2, "code_mapping": [mapping, None]}], ] for levels in cases: with self.subTest(levels=levels), self.assertRaises(ValueError): @@ -431,23 +443,6 @@ def test_single_target_mapping_is_allowed_anywhere(self): self.assertEqual(values.shape, (2, 1)) self.assertEqual(time.tolist(), [0.0, 1.0]) - def test_expansion_raises_when_it_would_duplicate_data(self): - mapping = self.MAPPING - cases: list[tuple[list[LevelConfig], Records]] = [ - ( - [{"ordering": "set", "tokens": 2, "code_mapping": [mapping, None]}], - {"tokens": [("401.9", "u")]}, - ), - ( - [CODES, {"ordering": "set", "tokens": 1, "code_mapping": mapping}], - {"tokens": ["401.9"], "children": [{"tokens": ["A"]}]}, - ), - ] - for levels, raw in cases: - proc = _fit(levels, raw) # fit only builds vocabularies - with self.subTest(raw=raw), self.assertRaises(ValueError): - proc.process(raw) - class TestResample(unittest.TestCase): @staticmethod @@ -847,6 +842,13 @@ def test_by_tokens_falls_back_to_all_tokens(self): # A: 0..2, B: 10..30, and unseen C: all tokens, 0..30. self.assertEqual(values.reshape(-1).tolist(), [0.5, 0.5, 0.5]) + def test_fit_checks_tokens_match_values(self): + level = self.level( + {"method": "zscore", "by_tokens": (0,)}, tokens=1, values=[1] + ) + with self.assertRaises(ValueError): + _fit([level], {"tokens": ["A", "B"], "values": [[1.0]]}) + def test_constant_divides_by_one(self): raw: Records = {"values": [[5.0], [5.0]]} proc = _fit([self.level({"method": "zscore"}, values=[1])], raw) @@ -1022,6 +1024,25 @@ def test_validate_rejects_broken_outputs(self): proc.validate(output) with self.assertRaises(ValueError): proc.validate((tokens, offsets, torch.tensor([0, 1])), batched=True) + with self.assertRaises(TypeError): + proc.validate((tokens.tolist(), offsets)) # type: ignore[arg-type] + + def test_validate_checks_shapes_values_and_time(self): + level: LevelConfig = {"ordering": "series", "tokens": 1, "values": [2]} + raw: Records = {"time": [0.0], "tokens": ["A"], "values": [[1.0, 2.0]]} + proc = _fit([level], raw) + tokens, values, time = proc.process(raw) + cases = { + "tokens [N, 1] for one namespace": (tokens[:, None], values, time), + "values [N, 1] for [2]": (tokens, values[:, :1], time), + "NaN value": (tokens, torch.full_like(values, NAN), time), + "time [N, 1]": (tokens, values, time[:, None]), + "infinite time": (tokens, values, torch.full_like(time, float("inf"))), + "record counts differ": (tokens, values, torch.zeros(2)), + } + for name, output in cases.items(): + with self.subTest(name), self.assertRaises(ValueError): + proc.validate(output) class TestStorage(unittest.TestCase): From 1e96214bee6a1e1b44d39fbefc9a5da0f7e35970 Mon Sep 17 00:00:00 2001 From: Logiquo Date: Fri, 30 Oct 2026 12:00:06 +0000 Subject: [PATCH 10/11] Address review: per-position normalization fallback and clearer docs - Normalization by token: each value position uses its token's statistic, or, where the token has none, the statistic over all tokens. A token whose training values were all NaN (or one value position of it) was left unnormalized, although the docs said it falls back like . The fallback is now one rule in _KeyedStatistics.lookup; imputation keeps its NaN for a key without a statistic, as documented. - A resampled level 0 without tokens or values is rejected: resampling drops its time, so it would have no output fields. - Docs: "merge" combines children into one set and keeps repeated children, as any set level does; only forward_fill and backward_fill fill tokens, linear leaves them. --- pyhealth/processors/sequence_processor_v2.py | 52 +++++++++++++------- tests/core/test_sequence_processor_v2.py | 40 +++++++++++++++ 2 files changed, 73 insertions(+), 19 deletions(-) diff --git a/pyhealth/processors/sequence_processor_v2.py b/pyhealth/processors/sequence_processor_v2.py index e5d0d55c6..b03ec4eb4 100644 --- a/pyhealth/processors/sequence_processor_v2.py +++ b/pyhealth/processors/sequence_processor_v2.py @@ -95,9 +95,10 @@ class Resample(TypedDict): not per token. - ``"min"`` / ``"max"`` / ``"avg"``: reduce the values. Only for a level with values and no tokens or children. - - ``"merge"``: union the records' children into one set. Only for - a level with children and no tokens or values, whose child - level is a ``set``. + - ``"merge"``: combine the records' children into one set; + repeated children are kept, as in any ``set`` level. Only for a + level with children and no tokens or values, whose child level + is a ``set``. A step without records gets ```` tokens, NaN values and no children; filling it is up to imputation. @@ -141,8 +142,10 @@ class Impute(TypedDict, total=False): edges as NaN. by_tokens: Token namespaces. Statistics are kept per token in these namespaces, and the time-based methods only use records - with the same token. None (the default) ignores tokens; the time- - based methods then also fill ```` tokens. + with the same token. None (the default) ignores tokens; + ``forward_fill`` and ``backward_fill`` then also fill ```` + tokens, while ``linear`` leaves them, as a token cannot be + interpolated. by_dims: For statistics: the axes of the value shape ``S`` to reduce over, indexed like ``values`` and ``spatial``. None (the default) reduces over all of them, ``()`` over none (one @@ -188,12 +191,13 @@ class Normalize(TypedDict, total=False): ``fit()`` from the raw training values (before resampling), with NaN skipped. A spread of 0 (a constant) divides by 1. by_tokens: Token namespaces; the statistics are kept per token in - these namespaces. A token without statistics uses the statistics - over all tokens, so its value keeps its measured size on the - population scale. That is ```` (any token not seen in - ``fit()``), ```` (a resampled step without records, whose - value is NaN unless imputed), and a token whose training values - were all NaN. None (the default) ignores tokens. + these namespaces. Each value position uses its token's statistic, + or, where the token has none, the statistic over all tokens, so + the value keeps its measured size on the population scale. That + covers ```` (any token not seen in ``fit()``), ```` (a + resampled step without records, whose value is NaN unless + imputed), and a token, or one value position of it, whose + training values were all NaN. None (the default) ignores tokens. by_dims: The axes of the value shape ``S`` to reduce over, indexed like ``values`` and ``spatial``. ``()`` (the default) keeps one statistic per value position, e.g. per channel of ``[N, F]`` @@ -475,10 +479,15 @@ def parse(cls, level: int, config: LevelConfig) -> "_InternalLevelConfig": f"Level {level}: resample method {method!r} reduces values, so the " "level needs values and no tokens or children (level 0 only)." ) + if level == 0 and not tokens and not values: + raise ValueError( + "Level 0: a resampled level drops its time, so it needs " + "'tokens' or 'values'." + ) if method == "merge" and (level == 0 or tokens or values): raise ValueError( - f"Level {level}: resample method 'merge' unions children, so the " - "level needs children and no tokens or values." + f"Level {level}: resample method 'merge' combines children, so " + "the level needs children and no tokens or values." ) resample = (float(step), method) @@ -679,7 +688,7 @@ class _KeyedStatistics: :meth:`finish` turns the running state into statistics, and :meth:`lookup` finds them for the records of a level in ``process()``. With ``pooled``, the statistics over all tokens are kept as well, under - the key None, for records whose key has none. + the key None, for the value positions where a key has none. """ def __init__( @@ -767,13 +776,18 @@ def lookup( ) -> Iterator[tuple[np.ndarray, dict[str, np.ndarray]]]: """Yields the mask of rows that share a key, and the key's statistics. - Rows whose key has no statistics are skipped, or, with ``pooled``, get - the statistics over all tokens. + Without ``pooled``, rows whose key has no statistics are skipped. With + ``pooled``, each value position uses the key's own statistic where it + has one, else the statistic over all tokens. """ + pooled = self._stats.get(None) if self.pooled else None for key, rows in _split_by_key(self._keys(tokens, n)): stats = self._stats.get(key) - if stats is None and self.pooled: - stats = self._stats.get(None) + if pooled is not None: + stats = pooled if stats is None else { + name: np.where(np.isnan(own), pooled[name], own) + for name, own in stats.items() + } if stats is not None: yield rows, stats @@ -1031,7 +1045,7 @@ def __init__(self, levels: Sequence[LevelConfig]): merge = config.resample is not None and config.resample[1] == "merge" if merge and self._config[level - 1].ordering != "set": raise ValueError( - f"Level {level}: resample method 'merge' unions the children's " + f"Level {level}: resample method 'merge' combines the children's " f"records, so level {level - 1} must have ordering 'set'." ) diff --git a/tests/core/test_sequence_processor_v2.py b/tests/core/test_sequence_processor_v2.py index 111d1fb31..46b6b9a9f 100644 --- a/tests/core/test_sequence_processor_v2.py +++ b/tests/core/test_sequence_processor_v2.py @@ -93,6 +93,8 @@ def test_rejects_bad_configurations(self): "resample": {**LAST, "method": "sum"}, } ], + # A resampled level drops its time; level 0 then has nothing left. + [{"ordering": "series", "resample": LAST}], [{"values": [1, 1], "spatial": [1, 0]}], [{"tokens": 1, "spatial": [False]}], # spatial without values [{"values": [1], "spatial": [False, True]}], # one flag per axis @@ -672,6 +674,17 @@ def test_statistics_come_from_raw_values(self): (values,) = _fit([level], raw).process(raw) self.assertEqual(values.tolist(), [[3.0], [5.0], [2.0]]) + def test_token_without_finite_values_keeps_its_nan(self): + # Unlike normalization, imputation does not fall back to all tokens: + # A keeps its NaN, which the chained "zero" step then fills. + train: Records = {"tokens": ["A", "B"], "values": [[NAN], [4.0]]} + step: Impute = {"method": "mean", "by_tokens": (0,)} + level = self.level(step, {"method": "zero"}, tokens=1, values=[1]) + _, values = _fit([level], train).process( + {"tokens": ["A", "B"], "values": [[NAN], [NAN]]} + ) + self.assertEqual(values.reshape(-1).tolist(), [0.0, 4.0]) + def test_fills_never_cross_groups(self): raw: Records = { "children": [ @@ -842,6 +855,33 @@ def test_by_tokens_falls_back_to_all_tokens(self): # A: 0..2, B: 10..30, and unseen C: all tokens, 0..30. self.assertEqual(values.reshape(-1).tolist(), [0.5, 0.5, 0.5]) + def test_by_tokens_falls_back_per_value_position(self): + # B never had a finite value: it uses the statistics over all tokens, + # like the unseen C. A never had one in channel 1: only that channel + # falls back. No token had one in channel 2: it is left as it is. + train: Records = { + "tokens": ["A", "A", "B", "D", "D"], + "values": [ + [0.0, NAN, NAN], + [2.0, NAN, NAN], + [NAN, NAN, NAN], + [10.0, 100.0, NAN], + [30.0, 300.0, NAN], + ], + } + level = self.level( + {"method": "minmax", "by_tokens": (0,)}, tokens=1, values=[3] + ) + proc = _fit([level], train) + _, values = proc.process( + {"tokens": ["A", "B", "C"], "values": [[1.0, 200.0, 7.0]] * 3} + ) + # Over all tokens: channel 0 is 0..30, channel 1 is 100..300. + torch.testing.assert_close( + values, + torch.tensor([[0.5, 0.5, 7.0], [1 / 30, 0.5, 7.0], [1 / 30, 0.5, 7.0]]), + ) + def test_fit_checks_tokens_match_values(self): level = self.level( {"method": "zscore", "by_tokens": (0,)}, tokens=1, values=[1] From d623252c76d65db9b31fbbcede3bf722d12cd3de Mon Sep 17 00:00:00 2001 From: Logiquo Date: Fri, 30 Oct 2026 12:36:19 +0000 Subject: [PATCH 11/11] Skip missing readings in min, max and avg resampling NaN in a value is a missing reading, as in the statistics and the missing indicators. The min, max and avg reductions propagated it, so one NaN reading in a step made the whole step NaN for that channel and threw away the other readings. They now skip NaN, so a step is NaN, and left to imputation, only where none of its readings has a value. --- pyhealth/processors/sequence_processor_v2.py | 21 ++++++++++++------ tests/core/test_sequence_processor_v2.py | 23 ++++++++++++++++++++ 2 files changed, 37 insertions(+), 7 deletions(-) diff --git a/pyhealth/processors/sequence_processor_v2.py b/pyhealth/processors/sequence_processor_v2.py index b03ec4eb4..d5e9259f6 100644 --- a/pyhealth/processors/sequence_processor_v2.py +++ b/pyhealth/processors/sequence_processor_v2.py @@ -93,8 +93,9 @@ class Resample(TypedDict): - ``"first"`` / ``"last"``: keep the first / last record, with its tokens, values and children. Both are over the whole step, not per token. - - ``"min"`` / ``"max"`` / ``"avg"``: reduce the values. Only for - a level with values and no tokens or children. + - ``"min"`` / ``"max"`` / ``"avg"``: reduce the values, skipping + NaN (a missing reading). Only for a level with values and no + tokens or children. - ``"merge"``: combine the records' children into one set; repeated children are kept, as in any ``set`` level. Only for a level with children and no tokens or values, whose child level @@ -651,12 +652,18 @@ def __call__( for i, kept in zip(filled.tolist(), keep.tolist()): grid_children[i] = children[kept] else: - reduce = {"min": np.minimum, "max": np.maximum, "avg": np.add}[self.method] - values = reduce.reduceat( - columns["values"].astype(np.float64), starts, axis=0 - ) + # NaN is a missing reading: it is skipped, so a step is NaN (for + # imputation) only where none of its readings has a value. + values = columns["values"].astype(np.float64) if self.method == "avg": - values /= (ends - starts + 1).reshape(-1, *[1] * (values.ndim - 1)) + seen = ~np.isnan(values) + total = np.add.reduceat(np.where(seen, values, 0.0), starts, axis=0) + count = np.add.reduceat(seen.astype(np.float64), starts, axis=0) + with np.errstate(invalid="ignore", divide="ignore"): + values = total / count + else: + reduce = np.fmin if self.method == "min" else np.fmax + values = reduce.reduceat(values, starts, axis=0) on_grid = np.full((n, *values.shape[1:]), np.nan, dtype=np.float32) on_grid[filled] = values resampled["values"] = on_grid diff --git a/tests/core/test_sequence_processor_v2.py b/tests/core/test_sequence_processor_v2.py index 46b6b9a9f..5ca2b16fb 100644 --- a/tests/core/test_sequence_processor_v2.py +++ b/tests/core/test_sequence_processor_v2.py @@ -536,6 +536,29 @@ def test_min_max_avg_reduce_values(self): proc = _fit([self.series(method, values=[2])], raw) self.assertEqual(proc.process(raw)[0].tolist(), expected) + def test_reductions_skip_missing_readings(self): + # Step 0: channel 0 has 1 and 3, channel 1 nothing. Step 1: one NaN + # reading next to a value in each channel. + raw: Records = { + "time": [0.0, 0.5, 1.0, 1.5], + "values": [[1.0, NAN], [3.0, NAN], [NAN, NAN], [5.0, 4.0]], + } + for method, expected in ( + ("min", [[1.0, 0.0], [5.0, 4.0]]), + ("max", [[3.0, 0.0], [5.0, 4.0]]), + ("avg", [[2.0, 0.0], [5.0, 4.0]]), + ): + with self.subTest(method=method): + level: LevelConfig = { + **self.series(method, values=[2]), + "impute": {"method": "zero"}, + "missing_indicator": True, + } + values, mask = _fit([level], raw).process(raw) + self.assertEqual(values.tolist(), expected) + # Only a step without any reading in a channel is missing. + self.assertEqual(mask.tolist(), [[False, True], [False, False]]) + def test_merge_unions_child_sets(self): raw: Records = { "time": [0.0, 0.5, 2.0],