Repository navigation
Towards Unified Sequence Representation I: Add SequenceProcessorV2 for packed hierarchical records - #1343
Towards Unified Sequence Representation I: Add SequenceProcessorV2 for packed hierarchical records#1343Logiquo wants to merge 8 commits into
Conversation
Test CoverageTotal: 74.29% (20804 / 28003 lines)
|
Lets callers check up front whether map() can return more than one target code, e.g. to reject a code mapping that would duplicate records. The icd-mappings backend always returns a single target; PyHealth tables are checked over their entries.
A single processor for nested codes, values and timestamps at any depth, emitted in CSR layout (payloads plus offsets) without padding or truncation. - Records: plain-dict columnar input (time, tokens, values, children). - LevelConfig: per-level ordering, token namespaces K, value shape S, code mapping, resampling and spatial axes; validated at construction. fit() only builds vocabularies. - Code mapping through CrossMap; multi-target mappings expand records and are only allowed on set levels without values. - Resampling of series levels onto a time grid with first, last, min, max, avg or merge; steps without records are left for imputation. - collate() and validate() for packed batches. - normalize, impute and missing_indicator are reserved for future work.
- Impute: a step {"method", "by_tokens", "by_dims"} or a chain of them,
run after resampling. Methods: zero (0.0 / <pad>), mean / min / max
(dataset statistics of the raw values, learned in fit(), per token in
by_tokens and reduced over the value axes in by_dims), forward_fill /
backward_fill / linear (along time, within one group, optionally only
from records with the same token). Gaps a step cannot fill keep NaN
and still raise unless a later step fills them.
- learns_statistics is True when an impute step uses statistics.
- Resampling and imputation now live in private _Resampler and _Imputer
classes; fit() and process() only call them.
- Normalize: {"method", "by_tokens", "by_dims"} with zscore or minmax,
run after resampling and imputation. Statistics are learned in fit()
from the raw values, per token in by_tokens and reduced over the value
axes in by_dims (default (): one statistic per value position). A token
without statistics (<unk>, <pad>, or only NaN in training) uses the
statistics over all tokens; a constant divides by 1.
- _KeyedStatistics holds per-token-key dataset statistics (count, mean,
M2 via Chan et al., min, max) and is shared by _Imputer and the new
_Normalizer.
- normalize is no longer reserved; Normalize is exported.
With missing_indicator, a level emits a "mask" role right after "values": a bool tensor shaped like the values, True where a value was NaN after resampling and before imputation (resampled steps without records included). It needs values and impute. collate() and validate() handle it like any other role. Nothing in the configuration is reserved any more.
…te() - stores_tensor now also covers a tuple of tensors with the same length for every sample; litdata stores such tuples natively, so SampleBuilder.encode writes them as tensors instead of pickling them. SequenceProcessorV2 declares stores_tensor = True. - A processor can define collate(items). collate_fn_dict_with_padding takes an optional mapping of per-field collators, checked before the (time, values) heuristic, and get_dataloader passes each processor's collate for its field. Other fields and loaders are unchanged.
forward_fill, backward_fill and linear reshaped the level's values with reshape(n, -1), which numpy cannot do when n is 0 (e.g. a sample whose stays have no lab records). Use the explicit value width instead, and test every impute method on such a level.
- Records, Resample, Impute, Normalize and LevelConfig get runnable '>>>' examples (all module doctests pass). - collate_fn_dict_with_padding gets an example and list / X | None annotations; its import block is sorted.
|
/review |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| # A code with several targets becomes several records, which only | ||
| # makes sense in an unordered level with no values to copy. | ||
| many = [m for m in mappers or () if m is not None and m.has_many_targets] | ||
| if many and (config.ordering != "set" or config.values): |
There was a problem hiding this comment.
[medium] Multi-target check in init misses cases that always fail at process()
The check only rejects multi-target code mappings when the level is not a set or has values. A set level with children (level > 0) or with several token namespaces (tokens > 1) passes construction, but _map_tokens (lines 1189-1199) raises ValueError on every sample that contains such a code. The LevelConfig docstring says the case is rejected by configuration, so users get a late runtime error instead of an early one. Fix: include level > 0 and config.tokens > 1 in the condition on line 1012, so these configurations fail in init with the existing message.
| 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 |
There was a problem hiding this comment.
[medium] Streaming mean/M2 update duplicates fit_normalization and TimeseriesProcessor logic
The Chan et al. parallel update here (lines 738-748) repeats the update in pyhealth/processors/utils.py:fit_normalization (lines 52-61), which already streams mean, M2, min and max over NaN-aware (N, F) arrays and produces the zscore/minmax centre and scale that _Normalizer re-derives. The forward_fill loop in _Imputer also repeats the forward fill in pyhealth/processors/timeseries_processor.py:144-151. The keyed (by_tokens/by_dims) version needs more than fit_normalization offers, so extend the shared helper to accept a key or per-key arrays and reuse it, rather than keeping a second copy of the algorithm. The PR does not mention these existing components.
| 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) |
There was a problem hiding this comment.
[low] time is stored as float32 before resampling
_read_time casts time to float32, and _Resampler bins it from that value. Absolute timestamps (for example Unix seconds around 1.7e9, with about 128 s float32 resolution) are rounded, so records can land in the wrong grid step or lose their order. Relative time in hours is fine. Either parse time as float64 for the grid computation and cast only the output, or document that time must be relative and small.
| @@ -0,0 +1,1540 @@ | |||
| """Unified processor for hierarchical packed records (CSR / jagged layout). | |||
|
|
|||
| See ``notes/Unified_Sequence_Processor.md`` (Section 7) for the design. | |||
There was a problem hiding this comment.
[low] Module docstring points to a design note that this PR does not add
The docstring refers to notes/Unified_Sequence_Processor.md (Section 7). That path is not among the PR's changed files, so readers of the repo cannot follow it. Remove the reference or move the relevant design text into the module or docs.
| for key, value in transformed.items(): | ||
| if key in tensor_fields: | ||
| if not isinstance(value, torch.Tensor): | ||
| is_tensors = isinstance(value, tuple) and all( |
There was a problem hiding this comment.
[low] Tuple-of-tensors storage does not enforce the same-length rule the docs state
The base-class and encode docs say a tuple stored with stores_tensor must have the same length in every sample, but encode only checks that every item is a tensor. A processor that varies the length will fail inside litdata with an unclear error rather than the ValueError the function raises for other type mismatches. Add a length check against the first sample, or soften the doc wording.
|
|
||
| # 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] |
There was a problem hiding this comment.
[low] schema() returns (level, role) keys, unlike the FeatureProcessor contract
The base class documents schema() as a tuple of semantic names. Existing consumers such as interpret/methods/deeplift.py:407 and integrated_gradients.py:322 test "mask" not in schema(), which is always True for V2 keys, so if a V2 field reaches those paths the check is silently wrong. The PR states that V2 is not consumable yet, so this is non-blocking, but consider raising NotImplementedError in the interpretability paths or documenting the divergence on the base class.
There was a problem hiding this comment.
AI Review
suggestion: need change
decision: human review
Thank you for this contribution. SequenceProcessorV2 gives one configurable way to represent nested codes, values and timestamps as packed records with no padding or truncation. It also adds the storage and batching support it needs. The tests are thorough, and coverage of the new processor file is 98%.
A maintainer will also take a look. The PR touches shared paths: SampleDataset.encode, get_dataloader, collate_fn_dict_with_padding and the FeatureProcessor base-class docs. Existing behaviour looks unchanged, but these paths are used by every dataset.
To get this merged:
- The constructor's multi-target check (around
sequence_processor_v2.py:1012) is incomplete. It lets through asetlevel that has children, and a level with several token namespaces. Such a config builds, then raises on every sample containing a multi-target code. Reject these cases in__init__, or make_map_tokenshandle them, and add tests. - The streaming mean/variance statistics (Chan et al.) and the forward-fill logic appear to duplicate what
pyhealth/processors/utils.py:fit_normalizationandTimeseriesProcessoralready provide. Please reuse those where you can. If they can't be reused, explain briefly why in the PR. - Coverage is below the 80% bar for
pyhealth/datasets/utils.py(45%) andpyhealth/processors/base_processor.py(about 80%, borderline). The whole repository is at 74%, which clears the 60% bar. Much of theutils.pygap is probably existing code. Please add tests for the new lines you added there, for examplecollate_fn_dict_with_paddingwithcollators.
Optional suggestions:
- Time is held as float32, which loses precision for absolute epoch timestamps. Consider float64, or document the limit.
- The module docstring cites
notes/Unified_Sequence_Processor.md, which isn't in the repo. Remove the reference or add the file. SampleBuilder.encodeaccepts tuples of tensors whose length varies between samples, although the docs require a fixed length. A stricter check, or a note in the docs, would help.
Reviewed commit 51f5c7e
This is the first PR of a stack towards one unified representation for sequential features. It adds
SequenceProcessorV2("sequence_v2"), a processor for nested codes, values and timestamps at any depth, stored as packed records (CSR / jagged layout) with no padding or truncation. It also makesSampleDatasetstore and batch such fields. Existing processors and models are unchanged. Model support comes in later PRs.Why
Today every combination of nesting, values and time gets its own processor (
SequenceProcessor,NestedSequenceProcessor,DeepNested*,StageNet*,Timeseries*, ...). Each one pads to rectangular shapes fixed duringfit(), so inputs longer than anything seen in training are truncated.SequenceProcessorV2describes all of them with one configuration, and its output keeps every record.What a field looks like
A task returns each field as
Records, a plain dict of columns.childrenholds the records one level down:{"time": [0.0, 30.0], "children": [{"tokens": ["A", "B"]}, {"tokens": ["C"]}]}The processor is configured per level, innermost first:
process()returns a tuple of tensors inschema()order: by level, then by role (tokens,values,mask,time,offsets). For example,((0, "tokens"), (0, "offsets"), (1, "time"))forvisits.offsetsat levellpartitions level-lrecords into level-l + 1records.Changes
pyhealth/processors/sequence_processor_v2.py(new)LevelConfig):ordering(set/sequence/series), the number of token namespacestokens(K) and the value shapevalues(S). These are fixed by the task, sofit()only learns vocabularies and statistics. The configuration is checked in__init__.<pad> = 0and<unk> = 1in each namespace.code_mappingmaps codes throughCrossMap. A mapping where one code has several targets expands records, so it is only allowed on asetlevel without values.serieslevel onto a time grid (Resample), withfirst,last,min,max,avgormerge(union of child sets). A step without records is left for imputation.Impute), a step or a chain of steps:zero;mean/min/max, dataset statistics of the raw values, per token (by_tokens) and reduced over value axes (by_dims);forward_fill/backward_fill/linearalong time, within one group, optionally only from the same token.Normalize):zscoreorminmax, with the same keys. A token without statistics falls back to the statistics over all tokens.boolmaskrole shaped likevalues, set where a value was missing before imputation.collate()packs a batch (concatenates payloads, shifts offsets, appends batch offsets[B + 1]), andvalidate()checks single-sample and batched outputs.fit()pass (count, mean and M2 via Chan et al., min, max) and pickle with the processor.pyhealth/medcode/cross_map.pyCrossMap.has_many_targets: whether some code maps to more than one target. Theicd-mappingsbackend never does.pyhealth/processors/base_processor.py,pyhealth/datasets/sample_dataset.pystores_tensor = Truenow also covers a tuple of tensors with the same length for every sample. litdata already stores such tuples natively, soSampleBuilder.encodewrites them as tensors and no longer pickles them.SequenceProcessorV2declares it.STORAGE_FORMATis unchanged: older caches still decode.pyhealth/datasets/utils.pycollate_fn_dict_with_padding(batch, collators=None): a field with a collator is batched only by it, before the(time, values)heuristic. Withoutcollators, behaviour is unchanged.get_dataloaderpasses each processor'scollate()(onlySequenceProcessorV2defines one) viafunctools.partial, so the collate function still pickles for worker processes.Docs:
docs/api/processors/pyhealth.processors.SequenceProcessorV2.rst, and exports frompyhealth.processors.Testing
tests/core/test_sequence_processor_v2.py: 76 tests covering:collate,validate, pickling,create_sample_datasetround trips;get_dataloaderbatching throughcollate().tests/core/test_medcode_icd_mappings.py:has_many_targets.SampleDataset,get_dataloaderand the unchangedTrainer, and trained successfully.Not in this PR
schema(),is_token()andsize()with their legacy meanings, so a V2 field cannot be passed to them yet. Adapters (packed_to_padded) and native jagged models are the next PRs in this stack.