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 9 commits into
Conversation
Test CoverageTotal: 74.68% (21033 / 28165 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}
There was a problem hiding this comment.
Previous review
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
- A one-to-many code_mapping is now rejected in __init__ unless the level is a bottom-level 'set' with one token namespace and no values. Before, a set level with children or several namespaces was accepted and then failed in process() on the first sample with such a code; that runtime check can no longer trigger and is removed. - SampleBuilder.encode checks that a stored tuple of tensors has one tensor per schema() entry, so every sample (in every worker) has the same structure for litdata. - SequenceProcessorV2._count drops an unreachable branch, and the module docstring no longer points to a design note outside the repo. - Tests cover the new rejections, collate_fn_dict_with_padding with collators, the tuple-length check, and the remaining validate, vocab, time and fit error paths.
|
Thanks for the review. Updated in d5420d8. Fixed
Intentionally not changed
|
|
/review |
| 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 and scale logic duplicates fit_normalization
_KeyedStatistics.accumulate re-implements the Chan et al. parallel update of mean and M2 that pyhealth/processors/utils.py fit_normalization (lines 52-61) already implements, and _Normalizer re-derives the zscore/minmax center and scale that fit_normalization (lines 68-78) already returns. This is a copy of an existing algorithm rather than a reuse. Fix: reuse fit_normalization per key, or extend it with a keyed variant, so the streaming update lives in one place.
| the statistics over all tokens. | ||
| """ | ||
| for key, rows in _split_by_key(self._keys(tokens, n)): | ||
| stats = self._stats.get(key) |
There was a problem hiding this comment.
[medium] Tokens with all-NaN training values do not fall back to pooled statistics
_KeyedStatistics.finish stores a statistics entry (with NaN values) for every key seen in fit(), including a token whose training values were all NaN (count 0). lookup() falls back to the pooled statistics only when the key is absent (stats is None), so such a token gets NaN statistics. _Normalizer then uses center 0 and spread 1 and leaves its values unnormalized. This contradicts the Normalize docstring, which says a token whose training values were all NaN uses the statistics over all tokens. No test covers this case. Fix: treat keys with count 0 as absent in finish() or lookup(), and add a test with an all-NaN token.
| 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) |
There was a problem hiding this comment.
[low] merge concatenates child tokens instead of taking a union
The Resample docstring says merge unions the records' children into one set, but grid_children[i].extend(kids) keeps duplicate child records (for example, the same code on two days in one step). Either document that merge concatenates, or deduplicate. The tests only use distinct tokens, so this is not exercised.
| if ( | ||
| tokens is not None | ||
| and step.by_tokens is None | ||
| and step.method != "linear" |
There was a problem hiding this comment.
[low] linear does not fill pad tokens, although the Impute docs say time-based methods do
The tokens-fill block is skipped for method 'linear' (step.method != 'linear'), but the Impute docstring says that without by_tokens 'the time-based methods then also fill tokens'. Align the docstring with the code, or the code with the docstring.
| ) | ||
| values = tuple(values) | ||
|
|
||
| if level == 0 and ordering != "series" and not tokens and not values: |
There was a problem hiding this comment.
[low] A level-0 series with only a resample produces an empty schema
The check accepts a level 0 with ordering 'series' and no tokens or values. If it also has resample (first or last), the time role is dropped and no roles remain, so schema() returns () and process() returns an empty tuple, which breaks collate() and validate(). Reject this configuration in parse().
|
|
||
| # 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() and spatial() overrides change the base FeatureProcessor contract
schema() returns (level, role) keys and spatial() returns a dict, but the base class defines schema() as a tuple of semantic names and spatial() as a tuple of bools. Existing models such as rnn.py:275 call schema().index('value'), and embedding.py:256 calls is_token(), so a V2 field fails in those models. The PR acknowledges that model support is out of scope, so this is non-blocking.
There was a problem hiding this comment.
AI Review
suggestion: need change
decision: human review
Thank you for this contribution. SequenceProcessorV2 gives PyHealth one configurable, padding-free packed representation for nested codes, values and time. The PR also adds tuple-of-tensor storage, a per-processor collate hook and CrossMap.has_many_targets. The tests are thorough, and the new processor file has 100% coverage.
A maintainer will take a look as well. The change touches shared core paths: SampleDataset.encode (the storage path) and get_dataloader / collate_fn_dict_with_padding (batching). Your description says behaviour is unchanged when no processor defines collate() or returns tuples. A maintainer still needs to confirm that.
I found no text aimed at reviewers.
Coverage: the whole repository is at 74.68%, above the 60% bar. The new processor file is at 100%, sample_dataset.py at 94%, cross_map.py at 97% and processors/init.py at 94%. Two modified files are below the 80% guideline: pyhealth/datasets/utils.py at 45% and base_processor.py at 80% (79.73%, just under). The low figure for utils.py comes mostly from code your PR doesn't touch. Please confirm that the lines you added are covered, in particular the get_dataloader collators branch.
To get this merged:
- Normalize fallback for all-NaN tokens. The
Normalizedocs say a token whose training values were all NaN uses the statistics over all tokens._KeyedStatistics.lookup()only falls back to the pooled statistics when the key is absent. A token with all-NaN training values has an entry whose mean, std, min and max are NaN. Such a token is therefore either left unnormalized or given the wrong scale. Either makelookupfall back in that case or change the docs, and add a test. - Reuse existing statistics code.
_KeyedStatisticsre-implements the streaming mean/variance and zscore/minmax logic. Please reusefit_normalizationinpyhealth/processors/utils.pywhere it fits. If it doesn't fit (per-token keys, the by_dims reduction), say why in a short comment.
Optional suggestions:
mergeis documented as a union of child sets, but duplicate child tokens appear to be concatenated. Either deduplicate or reword the docs.- The docs say
forward_fill/backward_fillfill<pad>tokens whenby_tokensis None, andlinearis documented to do the same. In the code,linearskips the token step. Please align the docs with the code. - A level-0
serieswithresampleand no tokens or values yields an empty schema. A configuration error would be clearer. - The
schema()andspatial()overrides change the base return types. They are marked as such already, but it is worth keeping that in mind for the later model PRs.
Reviewed commit d5420d8
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.