feat: handle Ensembl/symbol gene identifiers in each model's process_data - #418
Conversation
Review feedback on #418 (oriolpetithelical, mapping.py:275): no reason for it to be function-scoped. numpy is a hard dependency (`numpy>=2.1.3,<2.3` in pyproject) and is imported at module scope across the rest of helical, so the local import bought nothing -- it was an artefact of where the helper was first written. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
|
||
| def ensure_gene_symbols( | ||
| adata: AnnData, | ||
| gene_names: str = "index", |
There was a problem hiding this comment.
gene_names_locator maybe?
There was a problem hiding this comment.
Does not sound like a verb. We can do locate_gene_names. I do not have a hard opinion on this, WDYT?
There was a problem hiding this comment.
But why does it need to have a verb in it? This is just an identifier/locator of where to look for the genes in the dataset. But looking further in the code base, I see that gene_names is used accordingly. As such, keep it as it is
…data
Reconcile gene identifiers at each model's `process_data` -- the single choke
point shared by embed, fit, evaluate, eval and run_isp -- rather than in a
downstream consumer, which reached only two of those and needed a
hand-maintained model->namespace table kept in step by hand.
One shared set of primitives in `helical/utils/mapping.py`, replacing five inlined
copies of `startswith("ENS")`. That expression had three live bugs, each now
covered by a test:
- it matches real gene symbols (`ENSA`) and transcript/protein IDs (`ENST..`,
`ENSP..`), so anchored `^ENS[A-Z]{0,4}G\d{11}(\.\d+)?$` is used instead;
- `.all()`/`.any()` over the column cannot express a var index that is only
*mostly* Ensembl IDs, so detection is per entry;
- version suffixes were never stripped, and the mapping tables are keyed on bare
IDs, so every `ENSG..\.17` -- the GENCODE/CellRanger default -- resolved to None
and was dropped as "unmapped".
Per model, by what its vocabulary is keyed on:
- **scGPT, UCE, GenePT** (symbols): `ensure_gene_symbols` translates Ensembl IDs,
leaves existing symbols untouched, and collapses symbols claimed by more than
one gene -- 10616 of the 48698 symbol-bearing rows in the bundled table share a
gene_name, and a non-unique var index desyncs scGPT's count_matrix from its
gene_ids (#377) and raises InvalidIndexError in GenePT. Of a colliding set the
copy carrying the most counts wins; choosing positionally lets an all-zero
alt-scaffold copy displace the expressed one, after which the gene reads as
unexpressed with no error anywhere.
- **Geneformer, Tahoe** (Ensembl): `ensure_ensembl_ids` takes Ensembl input
**directly** instead of raising. Deliberately no symbol round trip: their
vocabularies contain genes with no gene symbol at all, so a round trip would
drop them. ENSG00000159239 -- in Geneformer's vocabulary, blank symbol in the
table -- is now tokenized rather than lost.
- **Nicheformer, Transcriptformer**: already branched on the identifier system
correctly; only their detector is swapped.
Two protections the removed guards provided by accident are restored explicitly,
and better, so the existing `test_ensembl_data_is_caught` passes **unmodified**:
- `require_vocabulary_overlap` rejects well-formed identifiers that are simply
from another annotation (mouse IDs against a human vocabulary) -- membership
rather than shape, so usable Ensembl input is accepted while unusable input
still fails loudly instead of tokenizing to nothing.
- `reject_null_identifiers` refuses literal "None"/"nan"/empty placeholders, which
mean an earlier mapping already failed. Matched exactly rather than by prefix,
so real genes like NANOS1 and NAT1 are unaffected.
Verified against real cached weights on CPU: an Ensembl-indexed AnnData that
previously failed with "No matching genes found between input data and scGPT gene
vocabulary" now tokenizes, for a plain Ensembl index, a versioned one, a mixed one
and a colliding one; Geneformer keeps the symbol-less in-vocabulary gene; mouse IDs
are rejected. 40 new tests for the primitives. Suite: 250 passed, 7 skipped, 1
pre-existing failure (transcriptformer gene mode, "Torch not compiled with CUDA
enabled", identical on the base commit); test_tahoe/helix_mrna/mamba2_mrna cannot
be collected on this host (flash_attn / mamba-ssm absent).
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
CI caught these; they cannot run on a host without flash_attn, so the Tahoe path was flagged as read-only-verified in the PR. Both tests asserted the behaviour this change deliberately replaces. - `test_process_data_raises_on_no_mapped_genes` patched `helical.models.tahoe.model.map_gene_symbols_to_ensembl_ids`, which no longer exists, to force an all-None ensembl_id column. Its *intent* -- nothing maps, so raise -- still holds, and UNKNOWN1/UNKNOWN2 are genuinely unmappable, so it now asserts that directly with no mock of an internal. That also stops it breaking the next time the helper changes. - `test_gene_mapping_ensembl_warning` asserted `match="ensemble ids"`: a column of Ensembl IDs with gene_names != "ensembl_id" was refused outright. That is exactly what this change does -- Tahoe's vocabulary *is* Ensembl-keyed, so those identifiers are now used as they are, with no symbol round trip. Rewritten to assert the new contract and renamed accordingly. It asserts on the reconciliation rather than the whole pipeline because the old test never reached the rest of `process_data` either (it always raised first) and the hand-built config in the fixture has none of what tokenization needs -- driving it further only produced an unrelated `KeyError: 'max_length'`. Unlike Geneformer's `test_ensembl_data_is_caught`, which passes unmodified because the protections it really guarded (vocabulary overlap, null sentinels) were restored explicitly, this one asserted the guard's *message* rather than a behaviour worth keeping, so re-baselining it is the honest fix. Both asserted behaviours were verified directly against `ensure_ensembl_ids` on the same data shape; the tests themselves can only be executed by CI. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Review feedback on #418 (oriolpetithelical, mapping.py:275): no reason for it to be function-scoped. numpy is a hard dependency (`numpy>=2.1.3,<2.3` in pyproject) and is imported at module scope across the rest of helical, so the local import bought nothing -- it was an artefact of where the helper was first written. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
New capability (every model's process_data accepts either gene-identifier system) plus new public helpers in helical/utils/mapping, so MINOR rather than PATCH. 3.0.4 was claimed by #417's squash merge. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Review feedback on #418 (oriolpetithelical, mapping.py:275): no reason for it to be function-scoped. numpy is a hard dependency (`numpy>=2.1.3,<2.3` in pyproject) and is imported at module scope across the rest of helical, so the local import bought nothing -- it was an artefact of where the helper was first written. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
0c8a9d7 to
e1a41b0
Compare
| out.var["original_gene_id"] = [ | ||
| value for value, kept in zip(identifiers, keep) if kept | ||
| ] | ||
| out.var_names = kept_names |
There was a problem hiding this comment.
This is a bit of surprise to me in the sense that even though the gene names are in a column like in line 374-375, we are also modifying the index. At least let's also document it in the docstring
There was a problem hiding this comment.
Also tied to this, I see an asymmetry between ensure_ensembl_ids versus ensure_gene_symbols. The one is always putting the outputs in a column (even though the values may be coming from an index) while the other one is putting the outputs in the index and optionally in the column
There was a problem hiding this comment.
Good point. This asymmetry is tributary of the asymmetry in consumers.
Added more information in the docstring.
Review feedback on #418 (oriolpetithelical, mapping.py:324). It was threaded through five functions purely for string interpolation, and it is redundant even for that: every caller's `process_data` logs "Processing data for <model>." immediately before calling these helpers, so the model already sits directly above the accounting line in the log stream, and an exception's traceback names the calling module. Messages are now model-agnostic. Removed from `ensure_gene_symbols`, `ensure_ensembl_ids`, `reject_null_identifiers`, `require_vocabulary_overlap` and `_log_accounting`; six call sites and two tests updated. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Review feedback on #418 (oriolpetithelical, mapping.py:324). It was threaded through five functions purely for string interpolation, and it is redundant even for that: every caller's `process_data` logs "Processing data for <model>." immediately before calling these helpers, so the model already sits directly above the accounting line in the log stream, and an exception's traceback names the calling module. Messages are now model-agnostic. Removed from `ensure_gene_symbols`, `ensure_ensembl_ids`, `reject_null_identifiers`, `require_vocabulary_overlap` and `_log_accounting`; six call sites and two tests updated. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Review feedback on #418 (oriolpetithelical, mapping.py:369): surprising that the index is modified even when the gene names live in a column. It is deliberate, and the docstring now says why: symbol-keyed models do not agree on which identifiers they read. scGPT reads `var[gene_names]`, but GenePT looks its embeddings up on `var_names` (`get_text_embeddings`) regardless of what `gene_names` was -- so leaving the index alone would silently match nothing there. Normalising both is what makes one result usable by any of them. The docstring now also states the rest of what the function writes: the named column is kept in step (callers run `ensure_rna_data_validity` first, which materialises a `var["index"]` from the pre-conversion index, and a lookup reading that stale column would match nothing), the pre-conversion identifiers are kept in `var["original_gene_id"]`, genes can be dropped so `n_vars` shrinks with `X` subset alongside, and the input is returned uncopied only when nothing needs converting. Every claim was verified by execution before being written down, and the two that were unpinned now have tests -- reverting the column sync turns both red. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Review feedback on #418 (oriolpetithelical, mapping.py:369): surprising that the index is modified even when the gene names live in a column. It is deliberate, and the docstring now says why: symbol-keyed models do not agree on which identifiers they read. scGPT reads `var[gene_names]`, but GenePT looks its embeddings up on `var_names` (`get_text_embeddings`) regardless of what `gene_names` was -- so leaving the index alone would silently match nothing there. Normalising both is what makes one result usable by any of them. The docstring now also states the rest of what the function writes: the named column is kept in step (callers run `ensure_rna_data_validity` first, which materialises a `var["index"]` from the pre-conversion index, and a lookup reading that stale column would match nothing), the pre-conversion identifiers are kept in `var["original_gene_id"]`, genes can be dropped so `n_vars` shrinks with `X` subset alongside, and the input is returned uncopied only when nothing needs converting. Every claim was verified by execution before being written down, and the two that were unpinned now have tests -- reverting the column sync turns both red. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…iffer Review feedback on #418 (oriolpetithelical): "I see an asymmetry between ensure_ensembl_ids versus ensure_gene_symbols. The one is always putting the outputs in a column (even though the values may be coming from an index) while the other one is putting the outputs in the index and optionally in the column." Correct observation; the asymmetry is deliberate, and the code said nothing about it. Each helper performs the **minimum mutation its consumers require**, and the consumers themselves are asymmetric: - Ensembl-keyed models read the column -- Geneformer's tokenizer reads `data.var.ensembl_id`, Tahoe reads `var[gene_id_key]` -- and neither uses `var_names` to identify a gene, so writing the column suffices. - Symbol-keyed models disagree with each other and between them cover both surfaces: GenePT and UCE read `var_names`, scGPT reads `var[gene_names]`. So `ensure_gene_symbols` has to normalise both. Leaving `var_names` alone on the Ensembl side is load-bearing rather than incidental: it keeps the caller's identifiers addressable for `id_to_gene` reverse lookups and caller-supplied gene lists, and rewriting the index would silently change the identifier system of anything reported downstream -- an ISP run would come back keyed on Ensembl IDs even when the caller supplied symbols. Documented on both functions so the pairing is discoverable from either. Every claim re-verified by execution: the column/index write behaviour, and each of the five consumer read-sites cited. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…iffer Review feedback on #418 (oriolpetithelical): "I see an asymmetry between ensure_ensembl_ids versus ensure_gene_symbols. The one is always putting the outputs in a column (even though the values may be coming from an index) while the other one is putting the outputs in the index and optionally in the column." Correct observation; the asymmetry is deliberate, and the code said nothing about it. Each helper performs the **minimum mutation its consumers require**, and the consumers themselves are asymmetric: - Ensembl-keyed models read the column -- Geneformer's tokenizer reads `data.var.ensembl_id`, Tahoe reads `var[gene_id_key]` -- and neither uses `var_names` to identify a gene, so writing the column suffices. - Symbol-keyed models disagree with each other and between them cover both surfaces: GenePT and UCE read `var_names`, scGPT reads `var[gene_names]`. So `ensure_gene_symbols` has to normalise both. Leaving `var_names` alone on the Ensembl side is load-bearing rather than incidental: it keeps the caller's identifiers addressable for `id_to_gene` reverse lookups and caller-supplied gene lists (helicalAI/bio-agent#1128, #1120), and rewriting the index would silently change the identifier system of anything reported downstream -- an ISP run would come back keyed on Ensembl IDs even when the caller supplied symbols. Documented on both functions so the pairing is discoverable from either. Every claim re-verified by execution: the column/index write behaviour, and each of the five consumer read-sites cited. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The helical repo should not carry references to another repository's tracker: they are unresolvable for anyone reading this code, and the bare numbers would render as links to unrelated helical issues. The reasoning each reference accompanied is kept in full -- only the identifiers are removed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2b00102 to
b81f25c
Compare
Reconciles gene identifiers inside each model's
process_data, so an AnnData indexed by Ensembl gene IDs works with the symbol-keyed models and vice versa. Stacked on #417 (GenePT's inverted guard), which is now merged.Approach
process_datais the single choke point shared byembed,fit,evaluate,eval/andrun_isp, so handling identifiers there covers all five at once. Per-model knowledge also stays in per-model code, rather than in a table that has to be kept in step by hand.One shared set of primitives, five inlined copies removed
startswith("ENS")appeared inline in Geneformer, Tahoe, Nicheformer, GenePT and Transcriptformer's dataloader. It had three live bugs, each now covered by a test:ENSA) and transcript/protein IDs (ENST…,ENSP…).all()/.any()can't express a mostly-Ensembl indexENSG….17— the GENCODE/CellRanger default — resolves toNoneand is droppedReplaced by
ENSEMBL_GENE_ID_PATTERN(anchored),is_ensembl_gene_id,ensembl_id_mask(per entry),strip_ensembl_version(Ensembl-matching values only — real symbols contain dots,AC000068.10).Per model, by what the vocabulary is keyed on
ensure_gene_symbols. Translates Ensembl IDs, leaves existing symbols alone so a mixed index keeps both, and collapses symbols claimed by more than one gene. That collapse is not optional: 10,616 of the 48,698 symbol-bearing rows in the bundled table share agene_name, and a non-unique var index desyncs scGPT'scount_matrixfrom itsgene_ids(scGPT.process_data desyncs gene_ids from count_matrix when gene names are non-unique #377) and raisesInvalidIndexErrorin GenePT. Of a colliding set the copy carrying the most counts wins — choosing positionally lets an all-zero alt-scaffold copy displace the expressed one, after which the gene reads as unexpressed with no error anywhere.ensure_ensembl_ids. Ensembl input is taken directly instead of raising, and deliberately not round-tripped through symbols: their vocabularies contain genes with no gene symbol at all, so a round trip drops them.ENSG00000159239— in Geneformer's vocabulary, blank symbol in the table — is now tokenized rather than lost.var_namesare left untouched, so reverse lookups and caller-supplied gene lists keep addressing the caller's own identifiers.The two helpers are deliberately asymmetric — one writes a column, the other the index — because their consumers are. Both docstrings now state the rule and the evidence.
Two protections restored explicitly
The guards being removed caught two things by accident, and an existing test encoded them. Rather than re-baseline that test, both are now checked directly — so
test_ensembl_data_is_caughtpasses unmodified:require_vocabulary_overlap— rejects well-formed identifiers from another annotation (mouse IDs against a human vocabulary). Membership rather than shape, so usable Ensembl input is accepted while unusable input still fails loudly instead of tokenizing to nothing. Neither Geneformer nor Tahoe had a zero-overlap check before.reject_null_identifiers— refuses literal"None"/"nan"/empty placeholders, which mean an earlier mapping step already failed. Matched exactly, not by prefix, so real genes likeNANOS1andNAT1are unaffected (the old check usedstartswith("None")).Verification
Against real cached weights on CPU:
ValueError: No matching genes foundIndexError/desync risk42 new tests for the primitives, including mixed indexes at several ratios,
ENSA, transcript/protein IDs, versioned IDs, and collapse-by-expression parametrised over which copy is expressed — positional selection passes one arrangement and fails the other, so a single fixture would prove nothing.Suite: 250 passed, 7 skipped, 1 failure (
test_transcriptformer … gene_mode,Torch not compiled with CUDA enabled, identical on the base commit).test_helix_mrnaandtest_mamba2_mrnacannot be collected on the dev host (mamba-ssmabsent);test_tahoeneedsflash_attnand is verified by CI, which caught two Tahoe tests encoding the previous contract — both since updated.No repo-wide reformatting: the diff is confined to the twelve files changed.