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..371da95bc --- /dev/null +++ b/docs/api/processors/pyhealth.processors.SequenceProcessorV2.rst @@ -0,0 +1,63 @@ +pyhealth.processors.SequenceProcessorV2 +======================================= + +Processor for hierarchical packed records: codes, values and timestamps at any +depth of nesting, stored without padding or truncation. + +The task returns each field as a ``Records``, a plain dict of columns. +``children`` holds the records one level down, and the innermost records +(level 0) have no ``children``. The processor emits a tuple of tensors in +``schema()`` order: by level, then by role (``tokens``, ``values``, ``mask``, +``time``, ``offsets``). ``offsets`` at level ``l`` partitions the level-``l`` records into +level-``l + 1`` records, in CSR layout. + +Each level's configuration declares its roles: ``tokens`` gives the number of +token namespaces ``K`` and ``values`` gives the value shape ``S``. ``fit()`` +only builds the vocabularies, and every record is checked against the +configuration. + +.. code-block:: python + + proc = SequenceProcessorV2(levels=[{"ordering": "set", "tokens": 1}, {"ordering": "sequence"}]) + samples = [{"visits": {"children": [{"tokens": ["A", "B"]}, {"tokens": ["C"]}]}}] + proc.fit(samples, "visits") + proc.schema() # ((0, 'tokens'), (0, 'offsets')) + proc.process({"children": [{"tokens": ["A", "B"]}, {"tokens": []}, {"tokens": ["C"]}]}) + # (tensor([2, 3, 4]), tensor([0, 2, 2, 3])) + +In a task schema the configuration is plain data: + +.. code-block:: python + + input_schema = { + "visits": ("sequence_v2", {"levels": [{"ordering": "set", "tokens": 1}, {"ordering": "sequence"}]}), + "vitals": ("sequence_v2", {"levels": [{"ordering": "series", "values": [2], "resample": {"step": 1.0, "method": "last"}}]}), + } + +Batches are packed with ``collate()``, which follows ``batch_schema()``: the +fields of ``schema()`` plus the batch offsets ``[B + 1]``. After resampling, +``impute`` fills NaN values (see ``Impute``) and ``normalize`` scales them (see +``Normalize``). With ``missing_indicator`` (which needs ``impute``), a level +also emits ``mask`` right after ``values``: a bool tensor of the same shape, +True where a value was missing before imputation. Values still containing NaN +or inf raise. + +.. autoclass:: pyhealth.processors.SequenceProcessorV2 + :members: + :undoc-members: + :show-inheritance: + +.. autoclass:: pyhealth.processors.Records + :members: + +.. autoclass:: pyhealth.processors.LevelConfig + :members: + +.. autoclass:: pyhealth.processors.Resample + :members: + +.. autoclass:: pyhealth.processors.Impute + :members: + +.. autoclass:: pyhealth.processors.Normalize + :members: diff --git a/pyhealth/datasets/sample_dataset.py b/pyhealth/datasets/sample_dataset.py index efcc47f32..b60702039 100644 --- a/pyhealth/datasets/sample_dataset.py +++ b/pyhealth/datasets/sample_dataset.py @@ -353,10 +353,11 @@ 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 - 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 + Fields whose processor declares ``stores_tensor`` stay tensors (or + 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: @@ -366,20 +367,34 @@ 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, 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(): if key in tensor_fields: - if not isinstance(value, torch.Tensor): - processor = self._input_processors.get(key) or self._output_processors.get(key) + 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): 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." + ) + # 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: diff --git a/pyhealth/datasets/utils.py b/pyhealth/datasets/utils.py index fa0acf5f7..55da18209 100644 --- a/pyhealth/datasets/utils.py +++ b/pyhealth/datasets/utils.py @@ -1,11 +1,13 @@ +import functools import hashlib import os import pickle +from collections.abc import Callable, Mapping from datetime import datetime from typing import Any, Dict, List, Optional, Tuple -import torch import litdata +import torch from dateutil.parser import parse as dateutil_parse from torch.nn.utils.rnn import pad_sequence from torch.utils.data import DataLoader @@ -250,16 +252,31 @@ def collate_fn_dict(batch: List[dict]) -> dict: return {key: [d[key] for d in batch] for key in batch[0]} -def collate_fn_dict_with_padding(batch: List[dict]) -> dict: +def collate_fn_dict_with_padding( + batch: list[dict], collators: Mapping[str, Callable] | None = None +) -> dict: """Collates a batch of data into a dictionary with padding for tensor values. Args: batch: List of dictionaries, where each dictionary represents a data sample. + collators: Optional collate function per key, given the key's values + in the batch. Keys with one are collated only by it. Returns: A dictionary where each key corresponds to a list of values from the batch. Tensor values are padded to the same shape. Tuples of (time, values) from temporal processors are collated separately. + + Examples: + >>> import torch + >>> batch = [{"x": torch.tensor([1, 2]), "id": "a"}, {"x": torch.tensor([3]), "id": "b"}] + >>> collate_fn_dict_with_padding(batch)["x"] + tensor([[1, 2], + [3, 0]]) + >>> collate_fn_dict_with_padding(batch)["id"] + ['a', 'b'] + >>> collate_fn_dict_with_padding(batch, collators={"x": len})["x"] + 2 """ collated = {} keys = batch[0].keys() @@ -267,8 +284,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 +422,25 @@ def get_dataloader( from .sample_dataset import SampleDataset if isinstance(dataset, SampleDataset): + # Fields whose processor defines collate() are batched by it. + processors = {**dataset.input_processors, **dataset.output_processors} + collators = { + key: processor.collate + for key, processor in processors.items() + if callable(getattr(processor, "collate", None)) + } + collate_fn = ( + functools.partial(collate_fn_dict_with_padding, collators=collators) + if collators + else collate_fn_dict_with_padding + ) return _SeededDataLoader( dataset, shuffle_samples=shuffle, seed=seed, batch_size=batch_size, generator=generator, - collate_fn=collate_fn_dict_with_padding, + collate_fn=collate_fn, ) if hasattr(dataset, "set_shuffle"): # another iterable dataset, e.g. litdata's dataset.set_shuffle(shuffle) 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/pyhealth/processors/__init__.py b/pyhealth/processors/__init__.py index b0738b05a..507fb0ffa 100644 --- a/pyhealth/processors/__init__.py +++ b/pyhealth/processors/__init__.py @@ -37,6 +37,14 @@ def get_processor(name: str): ) from .raw_processor import RawProcessor from .sequence_processor import SequenceProcessor +from .sequence_processor_v2 import ( + Impute, + LevelConfig, + Normalize, + Records, + Resample, + SequenceProcessorV2, +) from .signal_processor import SignalProcessor from .stagenet_processor import ( StageNetProcessor, @@ -76,6 +84,12 @@ def get_processor(name: str): "NestedSequenceProcessor", "RawProcessor", "SequenceProcessor", + "SequenceProcessorV2", + "Records", + "LevelConfig", + "Impute", + "Normalize", + "Resample", "SignalProcessor", "StageNetProcessor", "StageNetTensorProcessor", diff --git a/pyhealth/processors/base_processor.py b/pyhealth/processors/base_processor.py index 28915c89e..d3449869c 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 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; + :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 new file mode 100644 index 000000000..d5e9259f6 --- /dev/null +++ b/pyhealth/processors/sequence_processor_v2.py @@ -0,0 +1,1556 @@ +"""Unified processor for hierarchical packed records (CSR / jagged layout).""" + +from collections.abc import Iterable, Iterator, Mapping, Sequence +from dataclasses import dataclass +from itertools import accumulate, repeat +from typing import ( + TYPE_CHECKING, + Any, + Literal, + NamedTuple, + Required, + 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", "mask", "time", "offsets"] +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" +] +NormalizeMethod = Literal["zscore", "minmax"] + +# 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") + + +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. + + 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 + 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, 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 + is a ``set``. + + 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 + 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; + ``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 + 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. + + 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] + by_tokens: Sequence[int] | None + 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. 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]`` + 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. + + 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] + 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 the statistics of imputation and + normalization). + + 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 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. + 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: 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``. + + 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 + tokens: int + values: Sequence[int] + code_mapping: CodeMapping | Sequence[CodeMapping | None] | None + resample: Resample | None + spatial: Sequence[bool] | None + impute: Impute | Sequence[Impute] | None + normalize: Normalize | None + missing_indicator: bool + + +class _ImputeStep(NamedTuple): + """A checked :class:`Impute`.""" + + method: ImputeMethod + by_tokens: tuple[int, ...] | None + 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. + + 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 + impute: tuple[_ImputeStep, ...] + normalize: _NormalizeStep | None + missing_indicator: bool + + @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 + ) -> 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_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. + + ``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)}." + ) + + 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 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' combines 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) + + 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 = 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." + ) + 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)) + + 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 ( + 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, + values=values, + code_mapping=code_mapping, + resample=resample, + spatial=spatial, + impute=tuple(parsed), + normalize=normalize, + missing_indicator=missing_indicator, + ) + + +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: + # 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": + 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 + 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 the value positions where a 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. + + 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 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 + + +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 + (: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 = 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 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 statistics in self.statistics.values(): + statistics.finish() + + @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 + values = values.copy() + 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 + + # 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: + # 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) + 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} + + +@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``, ``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 + configuration; :meth:`fit` only builds the vocabularies and the statistics + of imputation and normalization. + + 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])) + """ + + # 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.") + 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 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] + 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 " + "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): + 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' combines the children's " + f"records, so level {level - 1} must have ordering 'set'." + ) + + 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 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( + 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 + ] + 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, + "mask": config.missing_indicator, + "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)) + 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: + """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 _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 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. + + 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] + 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]] = [] + 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) + ): + 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) + row.append(vocab[(targets or [code])[0]]) + ids.append(row) + if learners: + # Statistics are taken from the raw values, before + # resampling. + values = self._read_values(level, records) + for learner in learners: + learner.accumulate(values, ids) + children.extend(records.get("children", ())) + level_records = children + + 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, ...]: + """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]``, ``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 + 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) + 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] + if normalizer is not None: + normalizer(columns) + 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; impute them " + "(see Impute) or remove them in the task." + ) + 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 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``). + + 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 = { + "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}." + ) + 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 == "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): + 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_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_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().""" 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 new file mode 100644 index 000000000..5ca2b16fb --- /dev/null +++ b/tests/core/test_sequence_processor_v2.py @@ -0,0 +1,1175 @@ +import pickle +import unittest +from collections.abc import Sequence +from typing import Any +from unittest.mock import patch + +import torch + +from pyhealth.processors import ( + Impute, + 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_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): + 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"}, + } + ], + # 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 + [{"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_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(ValueError): + _fit([SCALARS]).vocab(0) # no tokens + 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"]}), + ([series], {"time": [[0.0]], "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): + 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")]}, + ) + 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_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): + 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]) + + +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_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], + "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, 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, + (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 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_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": [ + {"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_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"): + _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 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_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] + ) + 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) + 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 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"]) + 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) + 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): + 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))) + 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()