diff --git a/docs/api/processors.rst b/docs/api/processors.rst index f4f2996fb..fb0719457 100644 --- a/docs/api/processors.rst +++ b/docs/api/processors.rst @@ -24,6 +24,7 @@ Available Processors **Sequence Processors:** +- :doc:`SequenceProcessorV2 `: Unified processor for hierarchical packed records (codes, continuous values, and timestamps at any nesting depth) without padding or truncation - ``SequenceProcessor``: For categorical sequences (e.g., medical codes like diagnoses, procedures) - ``NestedSequenceProcessor``: For nested categorical sequences (e.g., drug recommendation with visit history) - ``NestedMultiHotProcessor``: For nested categorical sequences as per-visit multi-hot vectors (e.g., generative EHR models) @@ -312,9 +313,11 @@ How Processed Samples Are Stored ``set_task()`` and ``create_sample_dataset(..., in_memory=False)`` write processed samples to a disk cache (litdata). Each field is stored according to its processor: -- Processors that always return one ``torch.Tensor`` declare ``stores_tensor = True`` - (``sequence``, ``nested_sequence``, ``tensor``, ``timeseries``, the label - processors, ``image``, ``audio`` and others). Their outputs are stored as tensors. +- Processors that always return a ``torch.Tensor`` — or a tuple of tensors with + the same length for every sample — declare ``stores_tensor = True`` + (``sequence_v2``, ``sequence``, ``nested_sequence``, ``tensor``, ``timeseries``, + the label processors, ``image``, ``audio`` and others). Their outputs are + stored natively as tensors. - Every other field — ``raw``, ``text``, ``stagenet``, ``tuple_time_text``, ``graph``, ``temporal_timeseries``, ``time_image``, and keys that are not in the schema at all, such as ``patient_id`` or extra keys a task adds — is stored in @@ -329,8 +332,8 @@ already holds pickled processors in ``schema.pkl``: only load caches you created trust. A custom processor should set ``stores_tensor = True`` only if ``process()`` always -returns a single tensor; if such a field ever returns something else, writing the -cache fails with an error naming the field. +returns a single tensor or a fixed-length tuple of tensors; if such a field ever +returns something else, writing the cache fails with an error naming the field. In PyHealth 2.0.2 and earlier, non-tensor fields were handed to litdata as they were. litdata infers the layout of every sample from the first one it writes, so a @@ -338,14 +341,18 @@ value whose length, shape or type varied between samples failed to write, or was read back wrong. Caches written by those versions are rebuilt automatically (the storage layout is part of the cache key). -The default collate function treats any 2-tuple as a ``(time, values)`` pair of -tensors, so batch raw fields holding 2-tuples with your own ``collate_fn``. +Processors whose batches need custom packing (such as ``SequenceProcessorV2``) +define ``collate(items)``, which ``get_dataloader`` uses automatically for those +fields. For fields without a processor-level ``collate``, the default collate +function treats any 2-tuple as a ``(time, values)`` pair of tensors, so batch raw +fields holding 2-tuples with your own ``collate_fn``. Processor String Keys --------------------- Common string keys for automatic processor selection: +- ``"sequence_v2"``: For unified hierarchical packed records (codes, continuous values, and timestamps) - ``"temporal_timeseries"``: For time-series data with preserved timestamps (use in place of ``"timeseries"`` when building ``UnifiedMultimodalEmbeddingModel``) - ``"sequence"``: For categorical sequences (medical codes) - ``"nested_sequence"``: For nested categorical sequences (visit history) @@ -543,6 +550,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..e867a8339 --- /dev/null +++ b/docs/api/processors/pyhealth.processors.SequenceProcessorV2.rst @@ -0,0 +1,984 @@ +pyhealth.processors.SequenceProcessorV2 +======================================= + +:class:`~pyhealth.processors.SequenceProcessorV2` (registered as +``"sequence_v2"``) is PyHealth's unified feature processor for sequential, +hierarchical, and temporal clinical records. It represents categorical medical +codes, continuous physiological measurements, and timestamps at any depth of +nesting in a single lossless, padding-free Compressed Sparse Row (CSR / jagged) +layout. + +Quick Start +----------- + +A task supplies each feature as a :class:`~pyhealth.processors.Records` +dictionary of columns (``tokens``, ``values``, ``time``, and ``children``), and +configures the hierarchy from the innermost atomic level (level ``0``) outward +to level ``L``: + +.. code-block:: python + + from pyhealth.processors import Records, SequenceProcessorV2 + + # Level 0: unordered ICD-9 diagnosis codes within an encounter + # Level 1: chronological sequence of encounters for a patient + proc = SequenceProcessorV2( + levels=[ + {"ordering": "set", "tokens": 1}, + {"ordering": "sequence"}, + ] + ) + + train_samples = [ + { + "diagnoses": { + "children": [ + {"tokens": ["428.0", "401.9"]}, # Heart failure, hypertension + {"tokens": ["250.00"]}, # Type 2 diabetes + ] + } + } + ] + proc.fit(train_samples, "diagnoses") + proc.schema() + # ((0, 'tokens'), (0, 'offsets')) + + # Encounter 1 has no new diagnoses; "584.9" (acute kidney injury) was unseen in fit -> (1) + sample: Records = { + "children": [ + {"tokens": ["428.0", "401.9"]}, + {"tokens": []}, + {"tokens": ["250.00", "584.9"]}, + ] + } + tokens, offsets = proc.process(sample) + # tokens: tensor([2, 3, 4, 1]) — all 4 codes packed contiguously + # offsets: tensor([0, 2, 2, 4]) — encounter slices [0:2], [2:2] (empty), [2:4] + +In a :class:`~pyhealth.tasks.BaseTask`, configurations are plain data passed +via ``("sequence_v2", {"levels": [...]})`` or encapsulated in reusable +configuration subclasses: + +.. code-block:: python + + input_schema = { + "diagnoses": ( + "sequence_v2", + {"levels": [{"ordering": "set", "tokens": 1}, {"ordering": "sequence"}]}, + ), + "vitals": ( + "sequence_v2", + { + "levels": [ + { + "ordering": "series", + "values": [4], + "resample": {"step": 1.0, "method": "avg"}, + "impute": [{"method": "forward_fill"}, {"method": "zero"}], + "missing_indicator": True, + "normalize": {"method": "zscore"}, + } + ] + }, + ), + } + + +Why a Unified Sequence Representation? +-------------------------------------- + +Clinical Data Is Naturally Hierarchical and Multimodal +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Electronic health record (EHR) data spans multiple levels of clinical +granularity and combines discrete medical ontologies with continuous signals: + +- **1 level (flat events or signals):** A patient's chronological medication + administrations, a set of chronic conditions, or a stream of bedside ICU + vital signs. +- **2 levels (encounters or documents):** Diagnosis and procedure codes grouped + into outpatient visits, quantitative lab panels grouped by clinical + encounter, or word tokens grouped into timestamped progress notes. +- **3 levels (longitudinal hospital stays):** Medical codes grouped into + clinical encounters, and encounters grouped into distinct hospital + admissions across a patient's multi-year trajectory. + +Rather than treating each combination of nesting depth, timestamp presence, and +payload modality as an isolated preprocessing pipeline, +:class:`~pyhealth.processors.SequenceProcessorV2` formalizes hierarchy +directly through ``levels``. Adding timestamps to clinical visits or +introducing an outer hospital-admission grouping tier is a single declarative +change in ``levels`` rather than a rewrite of the dataset schema. + +One Clinical Semantic, One Canonical Representation +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Different model architectures often consume the exact same clinical data +through different tensor views. Consider a patient's longitudinal visit +history, where each visit records an unordered set of diagnosis codes: + +- Sequence and attention architectures often embed individual token indices, + either in native jagged layout or padded to ``[num_visits, max_codes]``. +- Bag-of-codes and generative architectures often represent each visit as a + binary multi-hot vector of shape ``[num_visits, vocab_size]``. + +Those two views describe the **exact same clinical phenomenon**—a sequence of +visits, each containing an unordered set of codes. Materializing model-specific +views during preprocessing fragments datasets into incompatible cached formats +and causes **premature information loss**: collapsing codes into a binary +multi-hot matrix during preprocessing irreversibly discards repeat code +occurrences, intra-group ordering, and code-level quantitative attributes (such +as medication dosages or lab values), preventing embedding-based models from +using the same processed dataset. + +:class:`~pyhealth.processors.SequenceProcessorV2` enforces the principle that +**identical clinical semantics always yield the same canonical, lossless +representation**. Raw events, values, timestamps, and group boundaries are +stored once; dense padded matrices, jagged nested tensors, or multi-hot +projections are derived on the model side. + +Zero Padding Waste and Zero Silent Truncation +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Clinical encounter sizes follow heavy-tailed distributions: routine outpatient +visits typically record 2–5 diagnosis codes, whereas complex critical-care +stays can record 80 or more. Enforcing fixed rectangular tensor shapes during +preprocessing introduces two severe flaws: + +1. **Silent test-time truncation:** Fixing ``max_codes_per_visit`` from the + training split silently truncates validation or test encounters whenever a + patient presents with more codes than the longest training visit. +2. **70–90% memory and cache waste:** Padding every routine visit up to the + dataset maximum fills disk caches, RAM, and host-to-device transfers + predominantly with ```` tokens (ID ``0``). + +By packing records contiguously and tracking group boundaries via 1D CSR index +``offsets``, :class:`~pyhealth.processors.SequenceProcessorV2` achieves **zero +padding and zero truncation** across preprocessing, native binary tensor disk +caching in :class:`~pyhealth.datasets.SampleDataset`, and batch collation in +:func:`~pyhealth.datasets.get_dataloader`. + + +Format Design Decisions +----------------------- + +Hierarchical ``levels`` and CSR ``offsets`` +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Data within a sample is organized into levels indexed from ``l = 0`` +(innermost atomic records, such as individual codes, lab measurements, vital +sign readings, or text tokens) to ``l = L`` (outermost grouping tier within a +sample, such as visits or hospital admissions). The ``levels`` sequence is +configured **innermost-first**: ``[level_0, level_1, ..., level_L]``. + +Whenever ``L > 0``, boundaries partitioning the ``N_l`` records at level ``l`` +into the ``N_{l+1}`` parent groups at level ``l + 1`` are stored as a 1D +``int64`` ``offsets`` tensor of shape ``[N_{l+1} + 1]`` obeying standard CSR +invariants: + +- **Contiguous slicing:** Group ``j`` at level ``l + 1`` owns the slice + ``records_l[offsets_l[j] : offsets_l[j + 1]]``. +- **Anchored endpoints:** ``offsets_l[0] == 0`` and ``offsets_l[-1] == N_l``. +- **Monotonicity and empty groups:** ``offsets_l[j] <= offsets_l[j + 1]``. + Equal adjacent offsets (``offsets_l[j] == offsets_l[j + 1]``) represent an + empty group—such as a clinical follow-up visit with zero new diagnosis + codes—without inserting dummy placeholder tokens. +- **Batch boundaries:** Within a single sample, the outermost level ``L`` has + no internal offsets. When batching ``B`` samples, :meth:`~pyhealth.processors.SequenceProcessorV2.collate` + concatenates all levels without padding, shifts each sample's internal + offsets by the cumulative record counts of preceding samples, and appends + outer batch offsets ``(L, "offsets")`` of shape ``[B + 1]``. + +Columnar Input (``Records``) and Positional Tensor Tuple Output +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +**Declared Columnar Input.** A task provides each feature as a +:class:`~pyhealth.processors.Records` dictionary holding columnar sequences at +one level (``tokens``, ``values``, ``time``, and ``children`` holding sub-records +one level down). Active roles are **declared in** +:class:`~pyhealth.processors.LevelConfig` **rather than inferred from data**: + +- Every sample—including empty patient histories (``N_l = 0``)—produces the + exact same number and dtypes of tensors. +- Missing declared keys or mismatched column lengths raise ``ValueError`` + immediately during :meth:`~pyhealth.processors.SequenceProcessorV2.process`. +- Undeclared extra keys in :class:`~pyhealth.processors.Records` are ignored, + making it inexpensive to iterate on ``input_schema`` without rewriting task + extraction logic. +- Timestamps are supplied as plain numbers in the task's chosen unit and + origin. Relative time (e.g., hours since ICU admission or days since index + diagnosis) is recommended because ``float32`` cannot represent absolute Unix + epoch seconds at sub-minute precision. + +**Positional Tensor Tuple and Processor Metadata.** Schema metadata lives once +on the processor (:meth:`~pyhealth.processors.SequenceProcessorV2.schema`, +:meth:`~pyhealth.processors.SequenceProcessorV2.batch_schema`, +:meth:`~pyhealth.processors.SequenceProcessorV2.ordering`, +:meth:`~pyhealth.processors.SequenceProcessorV2.vocab_size`, +:meth:`~pyhealth.processors.SequenceProcessorV2.value_dim`, +:meth:`~pyhealth.processors.SequenceProcessorV2.spatial`) and is never +duplicated inside cached samples. +:meth:`~pyhealth.processors.SequenceProcessorV2.process` returns a plain +``tuple[torch.Tensor, ...]`` ordered by **increasing level** (``0 .. L``), then +by **canonical role order**: + +.. list-table:: + :header-rows: 1 + :widths: 14 12 24 50 + + * - Role + - ``dtype`` + - Shape + - When Present + * - ``"tokens"`` + - ``int64`` + - ``[N_l]`` if ``K == 1``, else ``[N_l, K]`` + - ``tokens: K`` is configured at level ``l``. + * - ``"values"`` + - ``float32`` + - ``[N_l, *S_l]`` (scalars are ``[N_l, 1]``) + - ``values: S_l`` is configured at level ``l``. + * - ``"mask"`` + - ``bool`` + - ``[N_l, *S_l]`` + - ``missing_indicator: True`` (``True`` where ``NaN`` before imputation). + * - ``"time"`` + - ``float32`` + - ``[N_l]`` + - ``ordering == "series"`` and ``resample`` is ``None``. + * - ``"offsets"`` + - ``int64`` + - ``[N_{l+1} + 1]`` (or ``[B + 1]`` at ``L``) + - Every inner level ``l < L``; ``(L, "offsets")`` is appended in batches. + +Consumers can unpack the tuple positionally or bind tensors to ``(level, role)`` +keys via :meth:`~pyhealth.processors.SequenceProcessorV2.schema` (single +sample) or :meth:`~pyhealth.processors.SequenceProcessorV2.batch_schema` +(collated batch): + +.. code-block:: python + + fields = dict(zip(proc.schema(), proc.process(sample))) + tokens_0 = fields[(0, "tokens")] + offsets_0 = fields[(0, "offsets")] + + batched = proc.collate([proc.process(sample)]) + batch_fields = dict(zip(proc.batch_schema(), batched)) + batch_offsets = batch_fields[(proc.depth, "offsets")] # shape [B + 1] + +Because every sample emits a tuple of fixed length ``len(proc.schema())``, +:class:`~pyhealth.processors.SequenceProcessorV2` declares +``stores_tensor = True``: :class:`~pyhealth.datasets.SampleDataset` serializes +all tensors natively through ``litdata`` without falling back to per-sample +Python ``pickle``, while allowing every tensor's leading record count ``N_l`` +to vary freely across patients. + +Record ``ordering``: ``"set"`` vs. ``"sequence"`` vs. ``"series"`` +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Each level declares an ``ordering`` that captures the clinical ordering +semantics of its records and informs downstream models how to encode position: + +- **``"set"`` (unordered collection; no ``time`` tensor):** Record order carries + no clinical meaning—for example, ICD diagnosis codes or CPT procedure codes + assigned on a single hospital discharge summary (``{"428.0", "401.9"}`` is + clinically identical to ``{"401.9", "428.0"}``). Downstream models process + ``"set"`` levels with permutation-invariant pooling or self-attention + **without positional encodings**. +- **``"sequence"`` (discrete step order; no ``time`` tensor):** Discrete + step-by-step progression where order matters, either because inter-step + elapsed time is unmeasured (e.g., visit 1 -> visit 2 -> visit 3, or subword + tokens in a clinical note) or because a ``"series"`` level was placed onto a + regular time grid via ``resample``. Downstream models apply **discrete + index-based positional encodings** (indices ``0, 1, ..., n - 1``, such as + sinusoidal encodings, learned position embeddings, or RoPE). +- **``"series"`` (timestamped observations; ``time`` tensor required):** + Continuous or irregular time series where each record occurs at an explicit, + finite, non-decreasing timestamp (e.g., bedside ICU vital signs or serum lab + draws at ``t = 0.0, 1.5, 4.2`` hours). Downstream models apply + **continuous-time positional encodings** driven directly by ``time`` (such as + continuous sinusoidal embeddings, Time2Vec, or inter-event time-decay + :math:`\Delta t`). + +Categorical ``tokens`` and Multi-Namespace Ontologies (``K >= 1``) +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Configuring ``tokens: K`` declares ``K`` distinct categorical token namespaces +at that level. During :meth:`~pyhealth.processors.SequenceProcessorV2.fit`, +each namespace ``k in {0, ..., K - 1}`` builds an independent vocabulary +reserving `` = 0`` and `` = 1`` (any code unseen during training maps +to ```` at evaluation time, with no frequency cutoff). + +In clinical data, a single observation is frequently characterized by a tuple +of categorical attributes from different medical ontologies that must remain +bound to the same event. Concatenating attributes into a composite string (such +as ``"Glucose_mg/dL"``) causes combinatorial vocabulary explosion and prevents +models from learning shared representations across common tests, drugs, or +units. Setting ``tokens: K`` keeps independent vocabularies +(:meth:`~pyhealth.processors.SequenceProcessorV2.vocab_size` returns +``{level: (V_0, ..., V_{K-1})}``) so models can embed each coordinate +separately: + +1. **Quantitative Lab Measurements ``(test_code, unit)`` (``tokens: 2``):** + ``("Glucose", "mg/dL")`` vs. ``("Glucose", "mmol/L")``. Multi-center EHRs + often record the same analyte in different measurement units; separating + ``unit`` lets models calibrate the paired numeric ``values`` without + duplicating the lab test ontology. +2. **Medication Administrations ``(drug_code, route)`` (``tokens: 2``):** + ``("Insulin", "IV")`` vs. ``("Insulin", "SC")``. Administration route + fundamentally alters pharmacokinetics (onset, peak, duration) and clinical + acuity. +3. **Microbiology Cultures ``(specimen_site, organism)`` (``tokens: 2``):** + ``("Blood", "Staphylococcus aureus")`` vs. + ``("Sputum", "Pseudomonas aeruginosa")``. Isolating bacteria from sterile + blood versus non-sterile sputum carries distinct diagnostic implications for + bacteremia, sepsis, and antibiotic selection. +4. **Inpatient Diagnoses with Admission Status ``(diagnosis_code, poa_status)`` + (``tokens: 2``):** ``("J18.9", "POA_Y")`` vs. ``("J18.9", "POA_N")``. + "Present On Admission" (POA) flags separate pre-existing community-acquired + conditions from hospital-acquired complications. + +Continuous ``values`` and ``spatial`` Axis Semantics +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Configuring ``values: S`` (where ``S = [S_0, S_1, ...]``) declares a continuous +``float32`` payload of shape ``[N_l, *S]`` at level ``l``. Scalar measurements +always use ``values: [1]`` and emit shape ``[N_l, 1]`` (a flat 1D input list of +length ``N_l`` is also accepted for convenience). + +Because multi-dimensional clinical measurements mix unordered feature channels +with ordered continuous axes, ``spatial`` provides one boolean flag per axis of +``S`` (defaulting to ``(False,) * len(S)``): + +- **``False`` (unordered feature / channel axis):** Permuting channel order + does not alter the physical signal; downstream models apply dense projections + (``nn.Linear``). +- **``True`` (ordered spatial / high-frequency temporal axis):** Adjacent + elements share geometric or temporal neighborhood continuity; downstream + models apply convolutions (``nn.Conv1d``, ``nn.Conv2d``) or spatial patch + embeddings. + +.. list-table:: + :header-rows: 1 + :widths: 24 16 20 40 + + * - Clinical Modality + - ``values`` Shape + - ``spatial`` Flags + - Axis Semantics + * - Scalar lab measurement (e.g., Serum Potassium, 4.2 mEq/L) + - ``[1]`` + - ``[False]`` + - Single quantitative scalar attached to each lab observation. + * - Multichannel ICU vital panel (HR, SBP, DBP, MAP, RR, SpO2, Temp, GCS) + - ``[8]`` + - ``[False]`` + - 8 heterogeneous physiological variables; cross-channel interactions use ``nn.Linear`` rather than 1D convolutions across feature indices. + * - 12-lead ECG waveform strip (2 s at 500 Hz across leads I–III, aVR/L/F, V1–V6) + - ``[12, 1000]`` + - ``[False, True]`` + - Axis 0 (``12``, ``False``) indexes distinct anatomical leads; axis 1 (``1000``, ``True``) is a continuous waveform axis for ``nn.Conv1d`` filters. + * - Serial bedside Chest X-rays over consecutive ICU days + - ``[1, 224, 224]`` + - ``[False, True, True]`` + - Axis 0 (``1``, ``False``) is grayscale channel; axes 1–2 (``224 x 224``, ``True, True``) form a 2D pixel grid for ``nn.Conv2d`` or ViT patches. + +Built-In Clinical Preprocessing Pipeline +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Each level supports optional, declarative clinical preprocessing steps executed +in a strict, reproducible order during +:meth:`~pyhealth.processors.SequenceProcessorV2.process`: + +1. **Ontology translation (``code_mapping``):** Maps token namespaces through + :class:`~pyhealth.medcode.CrossMap` (e.g., ``("ICD9CM", "CCSCM")``) before + vocabulary lookup. Unmapped codes are preserved as-is. One-to-many ontology + mappings expand one input code into multiple records and are restricted to + bottom-level ``"set"`` configurations with ``tokens: 1`` and no ``values``. +2. **Temporal gridding (``resample``):** Bins a ``"series"`` level onto uniform + grid steps :math:`\lfloor t / \text{step} \rfloor` starting at :math:`t = 0` + using ``"first"``, ``"last"``, ``"min"``, ``"max"``, ``"avg"`` (skipping + ``NaN`` readings within a bin), or ``"merge"`` (taking the union of child + ``"set"`` records). Unobserved bins receive ```` tokens and ``NaN`` + values for downstream imputation. Output + :meth:`~pyhealth.processors.SequenceProcessorV2.ordering` becomes + ``"sequence"`` and ``time`` is dropped. +3. **Informative missingness tracking (``missing_indicator``):** When ``True``, + records a ``bool`` ``mask`` tensor (``[N_l, *S_l]``) immediately after + resampling and before imputation, marking every position that was originally + ``NaN`` or belonged to an unobserved resampled step. +4. **Chained missing-value imputation (``impute``):** Applies one + :class:`~pyhealth.processors.Impute` step or a list of steps in sequence + (``"zero"``, ``"mean"``, ``"min"``, ``"max"``, ``"forward_fill"``, + ``"backward_fill"``, ``"linear"``), optionally scoped per token key + (``by_tokens``) or reduced across selected value axes (``by_dims``). + Temporal fills operate strictly within each parent record's group and never + leak across visits or admissions. +5. **Feature scaling (``normalize``):** Applies :class:`~pyhealth.processors.Normalize` + (``"zscore"`` or ``"minmax"``) after imputation, scoped by ``by_tokens`` and + ``by_dims`` (default ``by_dims=()``, one scale per channel/value position). + Where a token lacks training statistics (````, ````, or all-``NaN`` + in training), each value position falls back to the pooled population + statistic rather than collapsing the measurement to zero. Any ``NaN`` or + ``inf`` remaining after the pipeline raises ``ValueError``. + +.. note:: + + Dataset statistics for ``impute`` (``mean``, ``min``, ``max``) and + ``normalize`` (``zscore``, ``minmax``) are learned in a single streaming + pass over **raw** training values (before resampling) during + :meth:`~pyhealth.processors.SequenceProcessorV2.fit`. Whenever statistics + are enabled (``proc.learns_statistics == True``), fit on the training split + via ``dataset.set_task(task, split=PatientSplit(...))`` so validation and + test patients do not influence preprocessing statistics. + + +Clinical Use Cases +------------------ + +1. Multi-Admission Longitudinal History with ICD-to-CCS Mapping (3 Levels) +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +A patient's longitudinal hospital trajectory spanning multiple admissions +(level ``2``), clinical encounters within each admission (level ``1``), and +unordered ICD-9-CM diagnosis codes mapped automatically to Clinical +Classifications Software (``CCSCM``) categories (level ``0``): + +.. code-block:: python + + from pyhealth.processors import Records, SequenceProcessorV2 + + adm_proc = SequenceProcessorV2( + levels=[ + { + "ordering": "set", + "tokens": 1, + "code_mapping": ("ICD9CM", "CCSCM"), + }, + {"ordering": "sequence"}, # Level 1: encounters within an admission + {"ordering": "sequence"}, # Level 2: hospital admissions + ] + ) + + patient_admissions: Records = { + "children": [ + { # Admission 0: ED intake (pneumonia, dyspnea) -> ICU transfer (resp. failure, severe sepsis) + "children": [ + {"tokens": ["486", "786.05"]}, + {"tokens": ["518.81", "995.92"]}, + ] + }, + { # Admission 1: Routine check (no new diagnoses) -> Cardiology consult (CHF, hypertension) + "children": [ + {"tokens": []}, + {"tokens": ["428.0", "401.9"]}, + ] + }, + ] + } + + adm_proc.fit([{"admissions": patient_admissions}], "admissions") + adm_proc.vocab(0) + # {'': 0, '': 1, '122': 2, '133': 3, '131': 4, '2': 5, '108': 6, '98': 7} + + adm_proc.schema() + # ((0, 'tokens'), (0, 'offsets'), (1, 'offsets')) + + tokens, enc_offsets, adm_offsets = adm_proc.process(patient_admissions) + # tokens: tensor([2, 3, 4, 5, 6, 7]) — 6 mapped CCS codes at level 0 + # enc_offsets: tensor([0, 2, 4, 4, 6]) — partitions 6 codes into 4 encounters (encounter 2 is empty) + # adm_offsets: tensor([0, 2, 4]) — partitions 4 encounters into 2 admissions + +2. Long-Format Quantitative Labs with ``(Test, Unit)`` Namespaces +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Irregularly sampled serum lab results in long format, binding ``(analyte, unit)`` +token pairs (``tokens: 2``) to quantitative scalar readings (``values: [1]``) +and relative hours since emergency department arrival. Missing numeric readings +are flagged in ``mask``, imputed per analyte (carry forward prior reading of +the same analyte -> analyte training mean -> ``0.0``), and z-score normalized +per ``(analyte, unit)`` pair: + +.. code-block:: python + + from pyhealth.processors import Records, SequenceProcessorV2 + + lab_proc = SequenceProcessorV2( + levels=[ + { + "ordering": "series", + "tokens": 2, + "values": [1], + "impute": [ + {"method": "forward_fill", "by_tokens": [0]}, + {"method": "mean", "by_tokens": [0]}, + {"method": "zero"}, + ], + "missing_indicator": True, + "normalize": {"method": "zscore", "by_tokens": [0, 1]}, + } + ] + ) + + train_labs = [ + { + "labs": { + "time": [0.0, 2.0, 4.0, 6.0], + "tokens": [ + ("Glucose", "mg/dL"), + ("Potassium", "mEq/L"), + ("Glucose", "mg/dL"), + ("Potassium", "mEq/L"), + ], + "values": [[100.0], [4.0], [120.0], [4.4]], + } + } + ] + lab_proc.fit(train_labs, "labs") + lab_proc.schema() + # ((0, 'tokens'), (0, 'values'), (0, 'mask'), (0, 'time')) + lab_proc.vocab_size() + # {0: (4, 4)} — independent vocabularies for analyte (ns 0) and unit (ns 1) + + nan = float("nan") + sample_labs: Records = { + "time": [0.5, 3.0, 5.5], + "tokens": [ + ("Glucose", "mg/dL"), + ("Glucose", "mg/dL"), # Pending / hemolyzed value -> carried forward from t=0.5 + ("Potassium", "mEq/L"), + ], + "values": [[110.0], [nan], [4.2]], + } + tokens, values, mask, time = lab_proc.process(sample_labs) + # tokens: shape [3, 2], int64 + # values: tensor([[0.], [0.], [0.]]) — normalized per (analyte, unit) + # mask: tensor([[False], [True], [False]]) + # time: tensor([0.5000, 3.0000, 5.5000]) + +3. Irregular Multichannel ICU Vital Signs Resampled Hourly +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Bedside Heart Rate (bpm) and Mean Arterial Pressure (mmHg) recorded at +irregular timestamps, aggregated into 1-hour averages +(``resample={"step": 1.0, "method": "avg"}``). Unobserved hours are marked in +``mask``, forward-filled from prior hours (with channel means as a leading-gap +fallback), and z-score scaled per vital channel: + +.. code-block:: python + + from pyhealth.processors import SequenceProcessorV2 + + vitals_proc = SequenceProcessorV2( + levels=[ + { + "ordering": "series", + "values": [2], + "resample": {"step": 1.0, "method": "avg"}, + "impute": [ + {"method": "forward_fill"}, + {"method": "mean", "by_dims": ()}, + {"method": "zero"}, + ], + "missing_indicator": True, + "normalize": {"method": "zscore"}, + } + ] + ) + + train_vitals = [ + { + "vitals": { + "time": [0.1, 0.6, 1.2, 2.5], + "values": [[80.0, 75.0], [84.0, 79.0], [90.0, 70.0], [100.0, 65.0]], + } + } + ] + vitals_proc.fit(train_vitals, "vitals") + vitals_proc.schema() # ((0, 'values'), (0, 'mask')) — time is replaced by 1-hour grid steps + vitals_proc.ordering() # {0: 'sequence'} + + # Two readings in hour 0 [0.0, 1.0), no readings in hour 1 [1.0, 2.0), one in hour 2 [2.0, 3.0) + values, mask = vitals_proc.process( + { + "time": [0.2, 0.8, 2.4], + "values": [[82.0, 76.0], [86.0, 78.0], [96.0, 68.0]], + } + ) + # values: shape [3, 2], float32 (hour 1 forward-filled from hour 0 average [84.0, 77.0], then scaled) + # mask: tensor([[False, False], [True, True], [False, False]]) + +4. Timestamped Clinical Progress Notes and Multi-Lead Waveforms +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Pre-tokenized clinical notes recorded across an ICU stay use 2 levels—word or +subword tokens as a ``"sequence"`` at level ``0``, grouped into timestamped +notes as a ``"series"`` at level ``1``—while multi-lead physiological waveforms +(such as 12-lead ECG segments sampled at 500 Hz for 2 seconds) use ``spatial`` +flags to distinguish anatomical leads from the temporal waveform axis: + +.. code-block:: python + + import numpy as np + from pyhealth.processors import Records, SequenceProcessorV2 + + # Pre-tokenized nursing/physician notes over elapsed ICU hours + notes_proc = SequenceProcessorV2( + levels=[ + {"ordering": "sequence", "tokens": 1}, # Level 0: tokens within a note + {"ordering": "series"}, # Level 1: timestamped notes + ] + ) + icu_notes: Records = { + "time": [0.0, 8.5], + "children": [ + {"tokens": ["chief", "complaint", "acute", "dyspnea"]}, + {"tokens": ["oxygen", "saturation", "improved", "on", "nasal", "cannula"]}, + ], + } + notes_proc.fit([{"notes": icu_notes}], "notes") + note_tokens, note_offsets, note_time = notes_proc.process(icu_notes) + # note_tokens: shape [10], int64 + # note_offsets: tensor([0, 4, 10]) + # note_time: tensor([0.0000, 8.5000]) + + # Serial 12-lead ECG strips (12 leads x 1000 samples per 2-second window) + ecg_proc = SequenceProcessorV2( + levels=[ + { + "ordering": "series", + "values": [12, 1000], + "spatial": [False, True], # 12 leads (channels), 1000 time ticks (spatial/temporal) + } + ] + ) + ecg_strips: Records = { + "time": [0.0, 6.0], + "values": np.zeros((2, 12, 1000), dtype=np.float32), + } + ecg_proc.fit([{"ecg": ecg_strips}], "ecg") + ecg_values, ecg_time = ecg_proc.process(ecg_strips) + # ecg_values: shape [2, 12, 1000], ecg_proc.spatial() == {0: (False, True)} + +5. Consuming Packed Batches in Downstream Models +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +From a single packed batch produced by +:meth:`~pyhealth.processors.SequenceProcessorV2.collate` (used automatically by +:func:`~pyhealth.datasets.get_dataloader`), downstream models can execute +natively on PyTorch 2.x jagged tensors without padding, materialize dense +rectangular tensors, or project unordered code sets into multi-hot vectors on +the fly: + +.. code-block:: python + + import torch + from torch.nn.utils.rnn import pad_sequence + + # Collate a batch of 2 patients (2 visits for patient 0, 1 visit for patient 1 -> 3 visits total) + p0 = proc.process({"children": [{"tokens": ["428.0", "401.9"]}, {"tokens": ["250.00"]}]}) + p1 = proc.process({"children": [{"tokens": ["428.0"]}]}) + tokens, visit_offsets, batch_offsets = proc.collate([p0, p1]) + + # Path A: Zero-padding PyTorch 2.x jagged NestedTensor (compatible with SDPA / FlashAttention) + jagged_visits = torch.nested.nested_tensor_from_jagged(tokens, visit_offsets) + + # Path B: Dense rectangular padded tensor [total_visits, max_codes_in_batch] + padded_visits = pad_sequence(jagged_visits.unbind(), batch_first=True, padding_value=0) + # tensor([[2, 3], + # [4, 0], + # [2, 0]]) + + # Path C: On-the-fly per-visit multi-hot matrix [total_visits, vocab_size] + num_visits = visit_offsets.numel() - 1 + vocab_len = proc.vocab_size()[0][0] + visit_ids = torch.repeat_interleave(torch.arange(num_visits), visit_offsets.diff()) + multihot_visits = torch.zeros(num_visits, vocab_len, dtype=torch.float32) + multihot_visits[visit_ids, tokens] = 1.0 + + +Mapping Common Clinical Data Structures +--------------------------------------- + +The table below maps common clinical data structures and modeling views to +their canonical :class:`~pyhealth.processors.Records` input, ``levels`` +configuration, output :meth:`~pyhealth.processors.SequenceProcessorV2.schema`, +and model-side consumption pattern: + +.. list-table:: + :header-rows: 1 + :widths: 20 24 24 16 16 + + * - Clinical Data Structure + - Raw ``Records`` Layout + - ``levels`` Configuration + - ``schema()`` & ``ordering()`` + - Model-Side View + * - **Ordered flat code sequence** (chronological diagnoses or procedures) + - ``{"tokens": ["428.0", "401.9"]}`` + - ``[{"ordering": "sequence", "tokens": 1}]`` + - ``((0, "tokens"),)``; ``{0: "sequence"}`` + - 1D ``int64`` token IDs ``[N_0]``; batch splits via ``(0, "offsets")``. + * - **Unordered patient code set** (bag-of-codes / flat multi-hot) + - ``{"tokens": ["428.0", "401.9"]}`` + - ``[{"ordering": "set", "tokens": 1}]`` + - ``((0, "tokens"),)``; ``{0: "set"}`` + - Embed as permutation-invariant set, or scatter ``tokens`` into ``[B, V]`` multi-hot. + * - **Visits of medical codes** (nested code sequences & per-visit multi-hot) + - ``{"children": [{"tokens": ["428.0", "401.9"]}, {"tokens": ["250.00"]}]}`` + - ``[{"ordering": "set", "tokens": 1}, {"ordering": "sequence"}]`` + - ``((0, "tokens"), (0, "offsets"))``; ``{0: "set", 1: "sequence"}`` + - Jagged ``(tokens, offsets)``, padded ``[V_tot, max_C]``, or scattered ``[V_tot, V]`` multi-hot. + * - **Admissions -> Visits -> Codes** (3-level longitudinal history) + - ``{"children": [{"children": [{"tokens": [...]}, ...]}, ...]}`` + - ``[{"ordering": "set", "tokens": 1}, {"ordering": "sequence"}, {"ordering": "sequence"}]`` + - ``((0, "tokens"), (0, "offsets"), (1, "offsets"))``; ``{0: "set", 1: "sequence", 2: "sequence"}`` + - ``(0, "offsets")`` groups codes into visits; ``(1, "offsets")`` groups visits into admissions. + * - **Timestamped visits of codes** (disease progression across irregular visits) + - ``{"time": [0.0, 30.0], "children": [{"tokens": [...]}, {"tokens": [...]}]}`` + - ``[{"ordering": "set", "tokens": 1}, {"ordering": "series"}]`` + - ``((0, "tokens"), (0, "offsets"), (1, "time"))``; ``{0: "set", 1: "series"}`` + - Codes at level 0 with visit boundaries ``(0, "offsets")`` and visit times ``(1, "time")``. + * - **Nested numeric measurements per visit** + - ``{"children": [{"values": [[1.2], [1.5]]}, {"values": [[0.9]]}]}`` + - ``[{"ordering": "sequence", "values": [F]}, {"ordering": "sequence"}]`` + - ``((0, "values"), (0, "offsets"))``; ``{0: "sequence", 1: "sequence"}`` + - Packed ``float32`` ``[N_0, F]`` partitioned into visits via ``(0, "offsets")``. + * - **Timestamped visits of numeric vectors** + - ``{"time": [0.0, 14.0], "children": [{"values": [[...]]}, {"values": [[...]]}]}`` + - ``[{"ordering": "sequence", "values": [F]}, {"ordering": "series"}]`` + - ``((0, "values"), (0, "offsets"), (1, "time"))``; ``{0: "sequence", 1: "series"}`` + - Numeric features at level 0 paired with visit-level timestamps ``(1, "time")``. + * - **Regularly gridded multichannel time series** + - ``{"values": [[80.0, 75.0], [82.0, 78.0]]}`` (or ``"series"`` + ``resample``) + - ``[{"ordering": "sequence", "values": [F]}]`` (or ``"series"`` with ``resample``) + - ``((0, "values"),)`` (+ ``(0, "mask")``); ``{0: "sequence"}`` + - Dense step sequence ``[N_0, F]`` per sample; batch boundaries in ``(0, "offsets")``. + * - **Irregular timestamped multichannel time series** + - ``{"time": [0.0, 1.5, 4.0], "values": [[80.0, 75.0], ...]}`` + - ``[{"ordering": "series", "values": [F]}]`` + - ``((0, "values"), (0, "time"))``; ``{0: "series"}`` + - Aligned ``float32`` ``values`` ``[N_0, F]`` and ``time`` ``[N_0]``. + * - **Long-format clinical events** (code / ``(code, unit)`` + value + timestamp) + - ``{"time": [...], "tokens": [("Glucose", "mg/dL"), ...], "values": [[105.0], ...]}`` + - ``[{"ordering": "series", "tokens": K, "values": [1]}]`` + - ``((0, "tokens"), (0, "values"), (0, "time"))``; ``{0: "series"}`` + - Co-indexed ontology IDs, quantitative values, and observation timestamps. + * - **Single tokenized clinical note** + - ``{"tokens": ["acute", "hypoxic", "respiratory", "failure"]}`` + - ``[{"ordering": "sequence", "tokens": 1}]`` + - ``((0, "tokens"),)``; ``{0: "sequence"}`` + - 1D ``int64`` token sequence (text arrives pre-split into tokens). + * - **Timestamped sequence of clinical notes** + - ``{"time": [0.0, 8.5], "children": [{"tokens": [...]}, {"tokens": [...]}]}`` + - ``[{"ordering": "sequence", "tokens": 1}, {"ordering": "series"}]`` + - ``((0, "tokens"), (0, "offsets"), (1, "time"))``; ``{0: "sequence", 1: "series"}`` + - Tokens partitioned into notes by ``(0, "offsets")``, aligned with ``(1, "time")``. + * - **Waveforms & serial medical images** (ECG strips, Chest X-ray series) + - ``{"time": [0.0, 24.0], "values": array}`` (e.g., ``[N_0, C, H, W]``) + - ``[{"ordering": "series", "values": [C, H, W], "spatial": [False, True, True]}]`` + - ``((0, "values"), (0, "time"))``; ``{0: "series"}`` + - Pre-loaded arrays with ``spatial()`` flagging channel vs. topological axes. + +Constructing ``Records`` from Common Python Data Structures +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Inside :meth:`BaseTask.__call__ `, common +Python collections and tabular EHR query outputs map directly into +:class:`~pyhealth.processors.Records` dictionaries: + +.. code-block:: python + + from pyhealth.processors import Records + + # 1. Flat list of codes: list[str] + flat_codes = ["428.0", "401.9", "250.00"] + rec_flat: Records = {"tokens": flat_codes} + + # 2. Nested list of visits: list[list[str]] + visits = [["428.0", "401.9"], [], ["250.00"]] + rec_visits: Records = {"children": [{"tokens": v} for v in visits]} + + # 3. 3-level admissions -> visits -> codes: list[list[list[str]]] + admissions = [[["486", "786.05"], ["518.81"]], [["428.0"]]] + rec_admissions: Records = { + "children": [ + {"children": [{"tokens": v} for v in adm]} + for adm in admissions + ] + } + + # 4. Timestamped visits of codes: list[tuple[float, list[str]]] + timed_visits = [(0.0, ["428.0", "401.9"]), (30.0, ["584.9"])] + rec_timed_visits: Records = { + "time": [t for t, _ in timed_visits], + "children": [{"tokens": codes} for _, codes in timed_visits], + } + + # 5. Irregular multichannel time series: (timestamps, values_2d) + times = [0.0, 1.5, 4.0] + vitals_matrix = [[80.0, 75.0], [84.0, 78.0], [78.0, 74.0]] + rec_vitals: Records = {"time": times, "values": vitals_matrix} + + # 6. Timestamped pre-tokenized clinical notes: list[tuple[float, list[str]]] + timed_notes = [(0.0, ["admit", "dyspnea"]), (12.0, ["stable", "on", "room", "air"])] + rec_notes: Records = { + "time": [t for t, _ in timed_notes], + "children": [{"tokens": words} for _, words in timed_notes], + } + + # 7. Long-format lab rows: list[tuple[float, str, str, float]] + lab_rows = [(0.0, "Glucose", "mg/dL", 108.0), (2.5, "Potassium", "mEq/L", 4.1)] + rec_labs: Records = { + "time": [t for t, _, _, _ in lab_rows], + "tokens": [(test, unit) for _, test, unit, _ in lab_rows], + "values": [[val] for _, _, _, val in lab_rows], + } + + +Defining Configuration Subclasses +--------------------------------- + +Users do not need to repeat raw ``levels=[...]`` dictionaries across task +schemas. :class:`~pyhealth.processors.SequenceProcessorV2` implements all +parsing, validation, vocabulary and statistics fitting, native disk storage, +and packed collation in one engine; **subclasses are pure configuration +presets** whose ``__init__`` passes a fixed (or parameterized) ``levels`` list +to ``super().__init__(levels=...)`` without overriding ``fit()``, ``process()``, +or ``collate()``. + +Registering configuration subclasses with +:func:`~pyhealth.processors.register_processor` makes them reusable by string +key or ``(key, kwargs)`` tuple in any task's ``input_schema``: + +.. code-block:: python + + from pyhealth.datasets import create_sample_dataset, get_dataloader + from pyhealth.processors import SequenceProcessorV2, register_processor + + + @register_processor("visit_codes") + class VisitCodesProcessor(SequenceProcessorV2): + """Unordered medical codes grouped into a sequence of clinical visits.""" + + def __init__(self, code_mapping: tuple[str, str] | None = None): + super().__init__( + levels=[ + {"ordering": "set", "tokens": 1, "code_mapping": code_mapping}, + {"ordering": "sequence"}, + ] + ) + + + @register_processor("admission_visit_codes") + class AdmissionVisitCodesProcessor(SequenceProcessorV2): + """Unordered codes grouped into visits, grouped into hospital admissions.""" + + def __init__(self, code_mapping: tuple[str, str] | None = None): + super().__init__( + levels=[ + {"ordering": "set", "tokens": 1, "code_mapping": code_mapping}, + {"ordering": "sequence"}, + {"ordering": "sequence"}, + ] + ) + + + @register_processor("timed_visit_codes") + class TimedVisitCodesProcessor(SequenceProcessorV2): + """Unordered medical codes grouped into timestamped clinical visits.""" + + def __init__(self): + super().__init__( + levels=[ + {"ordering": "set", "tokens": 1}, + {"ordering": "series"}, + ] + ) + + + @register_processor("hourly_vitals") + class HourlyVitalsProcessor(SequenceProcessorV2): + """Irregular multichannel vitals resampled hourly, imputed, and scaled.""" + + def __init__(self, num_channels: int, step_hours: float = 1.0): + super().__init__( + levels=[ + { + "ordering": "series", + "values": [num_channels], + "resample": {"step": step_hours, "method": "avg"}, + "impute": [ + {"method": "forward_fill"}, + {"method": "mean", "by_dims": ()}, + {"method": "zero"}, + ], + "missing_indicator": True, + "normalize": {"method": "zscore"}, + } + ] + ) + + + # Use the registered configuration subclasses directly in a task schema: + samples = [ + { + "patient_id": "p0", + "diagnoses": { + "children": [ + {"tokens": ["428.0", "401.9"]}, + {"tokens": ["584.9"]}, + ] + }, + "vitals": { + "time": [0.2, 1.4], + "values": [[82.0, 76.0, 98.0, 18.0], [88.0, 72.0, 97.0, 20.0]], + }, + "readmission": 1, + }, + { + "patient_id": "p1", + "diagnoses": {"children": [{"tokens": ["250.00"]}]}, + "vitals": { + "time": [0.5], + "values": [[74.0, 82.0, 99.0, 16.0]], + }, + "readmission": 0, + }, + ] + + dataset = create_sample_dataset( + samples=samples, + input_schema={ + "diagnoses": ("visit_codes", {"code_mapping": ("ICD9CM", "CCSCM")}), + "vitals": ("hourly_vitals", {"num_channels": 4}), + }, + output_schema={"readmission": "binary"}, + ) + batch = next(iter(get_dataloader(dataset, batch_size=2))) + # batch["diagnoses"] -> (tokens, visit_offsets, batch_offsets) + # batch["vitals"] -> (values, mask, batch_offsets) + + +API Reference +------------- + +.. 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()