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 11 commits into
Conversation
Test CoverageTotal: 74.68% (21041 / 28173 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 |
There was a problem hiding this comment.
Previous review
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
- Normalization by token: each value position uses its token's statistic, or, where the token has none, the statistic over all tokens. A token whose training values were all NaN (or one value position of it) was left unnormalized, although the docs said it falls back like <unk>. The fallback is now one rule in _KeyedStatistics.lookup; imputation keeps its NaN for a key without a statistic, as documented. - A resampled level 0 without tokens or values is rejected: resampling drops its time, so it would have no output fields. - Docs: "merge" combines children into one set and keeps repeated children, as any set level does; only forward_fill and backward_fill fill <pad> tokens, linear leaves them.
|
Thanks for the review. Updated in 1e96214. Fixed
Docs aligned with the code
Intentionally not changed
|
|
/review |
| for i, kept in zip(filled.tolist(), keep.tolist()): | ||
| grid_children[i] = children[kept] | ||
| else: | ||
| reduce = {"min": np.minimum, "max": np.maximum, "avg": np.add}[self.method] |
There was a problem hiding this comment.
[medium] avg/min/max resampling propagates NaN and discards valid readings
np.minimum, np.maximum and np.add propagate NaN, so one missing raw reading in a step makes the whole step NaN for that channel. process() accepts raw NaN (the missing-indicator tests rely on it), and the step is then marked missing and imputed, which throws away the other readings in the same step. Fix: reduce with NaN skipped (for example np.fmin/np.fmax for min and max, and a sum over finite values divided by the finite count for avg), and add a test with a NaN reading inside a step.
| 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.
[medium] Time is cast to float32, losing precision for absolute timestamps
Records.time is documented as numbers in the task's own unit, but it is cast to float32 (about 7 significant digits). Unix-second timestamps near 1.7e9 have a float32 spacing of 128 s, so a resample with step 1.0, or linear imputation, collapses distinct records into the same bins and produces wrong results without any error. Fix: keep time as float64 through _read_time and the resampler and emit float64 time (updating the dtype check in validate), or document that time must be relative and add a test with large offsets.
| 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] Chan mean/M2 merge duplicates the existing fit_normalization helper
This streamed update is the same Chan et al. merge as fit_normalization in pyhealth/processors/utils.py (lines 52-61), which tensor_processor.py and timeseries_processor.py already use. Fix: extract the merge arithmetic into a helper in processors/utils.py that works on per-key running state, and call it from _KeyedStatistics.accumulate, keeping the token-keyed and by_dims logic local to this module.
|
|
||
| # 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(), vocab_size() and spatial() overrides change the FeatureProcessor return types
The base class documents schema() as a tuple of semantic names and spatial() as a tuple of bools, and vocab_size() as an int. These overrides return (level, role) keys, dicts and dicts, respectively. The PR says V2 fields cannot reach legacy models yet, but interpretability code tests "mask" not in processor.schema(), which silently takes the non-mask path instead of failing. Consider raising a clear error from these methods for V2 processors, or documenting the divergence on the 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 processor for nested codes, values and timestamps stored as packed records with no padding or truncation, and the 76 tests and the end-to-end run are thorough. The new processor file has 100% coverage and the whole repository is at 74.68%, above the 60% bar. A maintainer will take a look as well, because the change modifies SampleDataset.encode (the disk-cache write path used by every dataset) and get_dataloader (the shared loader). Both changes look additive for existing processors, but they are core paths.
To get this merged:
- Fix the NaN handling in the
min/max/avgresampling. These reductions (np.minimum/np.maximum/np.addreduceat) propagate NaN. One NaN reading in a step makes the whole step NaN for that channel and throws away the valid readings. Use NaN-skipping reductions, for example a count of finite values foravg, so only steps with no finite readings stay NaN for imputation. Add a test for this. A test for empty groups would also help. - Fix the float32 time precision.
_read_timestores time as float32. Unix-second timestamps (about 1.7e9) lose precision at that size, so thefloor(time/step)binning and the linear interpolation can put records in the wrong step. Keep time as float64 during processing, or require and document relative time. The output dtype can stay float32 if the values are relative. - Reuse the existing streaming mean/variance code. The Chan et al. merge in
_KeyedStatistics.accumulateappears to duplicatefit_normalizationinpyhealth/processors/utils.py. Please reuse that helper, or explain why it cannot cover per-token, per-dimension, min and max statistics.
Optional suggestions:
pyhealth/datasets/utils.pyis at 45% coverage. The new collators branch and theget_dataloaderbranch are covered by your tests, but the 80% bar for modified files still applies. Please confirm the uncovered lines were already uncovered before this PR, or add tests for them.pyhealth/processors/base_processor.pyis at about 80%, which is borderline.schema(),spatial()andvocab_size()return types that differ from the FeatureProcessor contract. You note this in the PR, which is fine for now.
Reviewed commit 1e96214
NaN in a value is a missing reading, as in the statistics and the missing indicators. The min, max and avg reductions propagated it, so one NaN reading in a step made the whole step NaN for that channel and threw away the other readings. They now skip NaN, so a step is NaN, and left to imputation, only where none of its readings has a value.
|
Thanks for the review. Replies to each point:
|
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.