From 5a8e6e34d5193ddcf73890a2273dd7f1d3baf856 Mon Sep 17 00:00:00 2001 From: Alex Rubinsteyn Date: Mon, 5 Oct 2026 14:34:00 -0400 Subject: [PATCH 1/2] Make native mutation references portable and verified --- docs/native-references.md | 49 ++++++ requirements.txt | 3 +- tests/test_native_references.py | 244 +++++++++++++++++++++++++++++ vaxrank/cli/arg_parser.py | 4 + vaxrank/epitope_dataset.py | 34 +++- vaxrank/external_rescoring.py | 10 ++ vaxrank/mutant_protein_fragment.py | 9 ++ vaxrank/native_references.py | 229 +++++++++++++++++++++++++++ vaxrank/vaccine_antigen.py | 12 +- vaxrank/version.py | 2 +- 10 files changed, 584 insertions(+), 12 deletions(-) create mode 100644 docs/native-references.md create mode 100644 tests/test_native_references.py create mode 100644 vaxrank/native_references.py diff --git a/docs/native-references.md b/docs/native-references.md new file mode 100644 index 00000000..a8f401b4 --- /dev/null +++ b/docs/native-references.md @@ -0,0 +1,49 @@ +# Moving native mutation datasets + +Native epitope exports retain their evidence, scores, mutation fragments and +supporting transcript objects. Restoring transcript identity no longer opens an +annotation database; Vaxrank requires PyEnsembl 2.22.2 for that contract. + +When exporting a custom Genome, Vaxrank copies available annotation resources +into a sibling directory, for example: + +``` +result.tsv +result.tsv.references/ +``` + +Move both together. Each bundled resource has a relative path, byte size and +SHA-256 digest in the versioned `vaxrank.native_reference.v1` manifest. Identical +files share a copy. This can include complete GTF/transcript/protein/reference +DNA files, so the bundle size depends on the supplied custom reference. +Export does not download remote sources. Standard Ensembl references retain +their exact release/species identity; explicitly attached local reference DNA +is bundled too. + +Loading verifies available bundle files before rebinding paths. A checksum +conflict fails explicitly, even if an original annotation still exists. A +missing bundle file remains missing; Vaxrank does not silently substitute the +old source path or another annotation release. Gene/transcript/reference +identity and saved protein IDs survive relocation and re-save. + +Scoring and construct replay use the retained evidence without predictor calls +or annotation downloads. Even a table moved without its resource directory can +replay those measurements; annotation-dependent work needs the resources. +Availability is inspectable through `dataset.native_references.availability`. + +To prepare bundled local annotation indexes for full reports: + +```bash +vaxrank --input-epitopes result.tsv --index-native-references --output-dir reports +``` + +The equivalent Python call is `dataset.index_native_references()`. Indexing +verifies the resources again and acquires no data. Derived indexes are local to +the relocated bundle. Missing annotation resources produce an explicit error. + +Earlier native files retain their original reference identities and now restore +metadata lazily. Where the saved antigen already records protein IDs for the +same supporting transcripts, replay reuses them. Legacy files have no verified +resource manifest; Vaxrank cannot infer a replacement annotation or recover an +unavailable source. Re-export with the recorded reference available to bundle +it. diff --git a/requirements.txt b/requirements.txt index fab30fa2..4bd8f1a7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,7 +8,8 @@ datacache>=1.14.0,<2.0.0 # Public data acquisition and original-read extraction for the osteosarc corpus. osteosarc>=0.14.4,<0.15 pandas>=2.1.4,<4.0.0 -pyensembl>=2.16.0,<3.0.0 +# Metadata-only Transcript/Gene restore defers annotation DB access (#563). +pyensembl>=2.22.2,<3.0.0 # 10.4.1 ranks structural candidate sets that list themselves as candidates # without recursing, which DNA fallback relies on when comparing fusion outcomes. # Varcode 10 also supplies candidate-local SV sequence-change checks, including diff --git a/tests/test_native_references.py b/tests/test_native_references.py new file mode 100644 index 00000000..8aea6afd --- /dev/null +++ b/tests/test_native_references.py @@ -0,0 +1,244 @@ +"""Actual mutation evidence and annotations survive relocation independently.""" + +from dataclasses import replace +import json +from pathlib import Path +import shutil +import subprocess +import sys + +import pandas as pd +import pytest +from pyensembl.download_cache import MissingLocalFile + +from vaxrank.cli.entry_point import run_cli +from vaxrank.epitope_config import EpitopeConfig +from vaxrank.epitope_dataset import EpitopeDataset, DATASET_METADATA, read_table +from vaxrank.epitope_dsl import attach_per_allele_scores +from vaxrank.native_references import NativeReferences, REFERENCE_FIELD +from vaxrank.mutant_protein_fragment import MutantProteinFragment +from vaxrank.native_serialization import from_native_json, to_native_json +from vaxrank.vaccine_antigen import VaccineAntigen +from .test_core_logic_config import _make_epitope +from .test_epitope_dataset import forbidden +from .test_isovar_support_contract import MUTANT, reconstruct, reference as _support_reference + + +@pytest.fixture(params=["+", "-"]) +def reference(tmp_path, request): + return _support_reference.__wrapped__(tmp_path, request) + + +def native_dataset(tmp_path, reference): + result = reconstruct(tmp_path, reference, [(MUTANT, [(len(MUTANT), "M")], 0)] * 2) + fragment = MutantProteinFragment.from_isovar_result(result) + epitope = replace(_make_epitope(fragment.amino_acids[1:10], ic50=100., wt_ic50=1000., + source_sequence=fragment.amino_acids, offset=1), prediction_id="occurrence") + config = EpitopeConfig(score_expr="1 / affinity.value", min_epitope_score=0) + dataset = EpitopeDataset.from_predictions([epitope], config=config) + dataset.epitopes = tuple(attach_per_allele_scores(dataset.epitopes, config, + topiary_df=dataset.scoring_frame())) + dataset.mutation_fragments["occurrence"] = fragment + dataset.antigens["occurrence"] = VaccineAntigen.from_mutant_protein_fragment(fragment) + path = tmp_path / "archive" / "native.tsv" + dataset.save(path) + return dataset, path + + +def delete_original_reference(tmp_path): + for name in ("ref.gtf", "cdna.fa", "protein.fa"): + (tmp_path / name).unlink() + shutil.rmtree(tmp_path / "cache") + + +def test_mutation_native_bundle_relocates_without_old_sources_or_predictors(tmp_path, reference, monkeypatch): + import mhctools.cli + import pyensembl.download_cache + before, path = native_dataset(tmp_path, reference) + shutil.move(str(path.parent), str(tmp_path / "relocated")) + path = tmp_path / "relocated" / "native.tsv" + delete_original_reference(tmp_path) + monkeypatch.setattr(mhctools.cli, "predictors_from_args", forbidden) + monkeypatch.setattr(pyensembl.download_cache.DownloadCache, "_fetch", forbidden) + loaded = EpitopeDataset.load(path) + assert all(item["available"] for item in loaded.native_references.availability) + fragment = loaded.mutation_fragments["occurrence"] + transcript, = fragment.supporting_reference_transcripts + assert transcript.id == "t" and transcript.gene_id == "g" + assert transcript.genome.reference_name == before.mutation_fragments["occurrence"].variant.reference_name + assert (transcript.genome.annotation_name, transcript.genome.annotation_version) == ("test", 1) + pd.testing.assert_frame_equal(loaded.result.df[before.result.df.columns], before.result.df) + rescored = attach_per_allele_scores(loaded.epitopes, loaded.config, topiary_df=loaded.scoring_frame()) + assert [e.per_allele_scores for e in rescored] == [e.per_allele_scores for e in before.epitopes] + # Explicitly prepare local annotation indexes, never a remote/default release. + loaded.index_native_references() + assert transcript.gene.name == "G" and transcript.protein_sequence == "MAAAAKGGGG" + assert fragment.predicted_effect() is not None + again = tmp_path / "second" / "again.tsv" + loaded.save(again) + second = EpitopeDataset.load(again) + assert second.antigens == loaded.antigens + assert second.mutation_fragments["occurrence"].supporting_reference_transcripts[0].id == "t" + # A fresh process must resolve the moved relative resources independently. + code = """ +import sys +from pyensembl.download_cache import DownloadCache +from vaxrank.epitope_dataset import EpitopeDataset +DownloadCache._fetch = lambda *a, **k: (_ for _ in ()).throw(AssertionError('no downloads')) +d = EpitopeDataset.load(sys.argv[1]) +assert d.mutation_fragments['occurrence'].supporting_reference_transcripts[0].id == 't' +d.index_native_references() +assert d.mutation_fragments['occurrence'].supporting_reference_transcripts[0].gene.name == 'G' +""" + subprocess.run([sys.executable, "-c", code, str(again)], check=True) + + +def test_missing_bundle_allows_evidence_replay_and_reports_unavailable_annotation(tmp_path, reference): + _, path = native_dataset(tmp_path, reference) + delete_original_reference(tmp_path) + shutil.rmtree(Path(str(path) + ".references")) + loaded = EpitopeDataset.load(path) + assert not any(item["available"] for item in loaded.native_references.availability) + assert loaded.epitopes[0].per_allele_scores == {"HLA-A*02:01": .01} + assert loaded.mutation_fragments["occurrence"].supporting_reference_transcripts[0].id == "t" + with pytest.raises(ValueError, match="resources are unavailable"): + loaded.index_native_references() + with pytest.raises(MissingLocalFile, match="references"): + loaded.mutation_fragments["occurrence"].supporting_reference_transcripts[0].gene + + +def test_missing_reference_identity_survives_cli_resave(tmp_path, reference, monkeypatch): + import mhctools.cli + import pyensembl.download_cache + _, path = native_dataset(tmp_path, reference) + expected = EpitopeDataset.load(path).native_references.availability + delete_original_reference(tmp_path) + shutil.rmtree(Path(str(path) + ".references")) + monkeypatch.setattr(mhctools.cli, "predictors_from_args", forbidden) + monkeypatch.setattr(pyensembl.download_cache.DownloadCache, "_fetch", forbidden) + second = tmp_path / "replay" / "native.tsv" + run_cli(["--input-epitopes", str(path), "--output-epitopes", str(second), + "--vaccine-peptide-length", "10", "--no-processing-aware-annotation"]) + loaded = EpitopeDataset.load(second) + def identity(resources): + return {(r['field'], r['index'], r['sha256'], r['size']) for r in resources} + assert identity(loaded.native_references.availability) == identity(expected) + assert not any(r['available'] for r in loaded.native_references.availability) + assert loaded.antigens + with pytest.raises(ValueError, match="resources are unavailable"): + loaded.index_native_references() + + +def test_legacy_archive_reuses_recorded_protein_ids_without_annotation(tmp_path, reference, monkeypatch): + import pyensembl.download_cache + _, path = native_dataset(tmp_path, reference) + table = read_table(path) + payload = table.extra[DATASET_METADATA] + def legacy(value): + if isinstance(value, list): + return [legacy(v) for v in value] + if isinstance(value, dict): + return {k: legacy(v) for k, v in value.items() + if k not in (REFERENCE_FIELD, 'supporting_reference_protein_ids')} + return value + for key, value in payload['mutation_fragments'].items(): + payload['mutation_fragments'][key] = json.dumps(legacy(json.loads(value))) + # Write through the public table codec, preserving its typed evidence. + table.to_tsv(path) + delete_original_reference(tmp_path) + shutil.rmtree(Path(str(path) + ".references")) + monkeypatch.setattr(pyensembl.download_cache.DownloadCache, "_fetch", forbidden) + loaded = EpitopeDataset.load(path) + assert loaded.mutation_fragments['occurrence'].supporting_reference_protein_ids == ('p',) + second = tmp_path / 'legacy-replay.tsv' + run_cli(['--input-epitopes', str(path), '--output-epitopes', str(second), + '--vaccine-peptide-length', '10', '--no-processing-aware-annotation']) + assert EpitopeDataset.load(second).antigens + + +def test_changed_bundle_is_rejected_even_if_original_reference_still_exists(tmp_path, reference): + _, path = native_dataset(tmp_path, reference) + source = next(Path(str(path) + ".references").glob("*.gtf")) + source.write_text(source.read_text().replace("gene_name", "gene_fame")) # Same byte length. + with pytest.raises(ValueError, match="checksum mismatch"): + EpitopeDataset.load(path) + + +def test_reference_change_after_loading_is_rejected_on_resave(tmp_path, reference): + _, path = native_dataset(tmp_path, reference) + loaded = EpitopeDataset.load(path) + source = next(Path(str(path) + ".references").glob("*.gtf")) + source.write_text(source.read_text().replace("gene_name", "gene_fame")) + with pytest.raises(ValueError, match="checksum mismatch"): + loaded.save(tmp_path / "changed.tsv") + + +def test_ensembl_release_preserves_identity_and_bundles_attached_dna(tmp_path, monkeypatch): + from pyensembl import EnsemblRelease + from varcode import Variant + import pyensembl.download_cache + monkeypatch.setattr(pyensembl.download_cache.DownloadCache, "_fetch", forbidden) + dna = tmp_path / "dna.fa" + dna.write_text(">1\nACGTACGTACGT\n") + genome = EnsemblRelease(104, species="mouse", genome_fasta=str(dna)) + dataset = EpitopeDataset.from_predictions([_make_epitope("SIINFEKL")]) + dataset.direct_sources = [{'variant': to_native_json(Variant('1', 2, 'C', 'T', ensembl=genome)), + 'properties': {}}] + path = tmp_path / "archive" / "native.tsv" + dataset.save(path) + shutil.move(str(path.parent), str(tmp_path / 'moved')) + dna.unlink() + loaded = EpitopeDataset.load(tmp_path / 'moved' / 'native.tsv') + variant = from_native_json(loaded.direct_sources[0]['variant'], Variant) + assert variant.genome.release == 104 and variant.genome.species.latin_name == 'mus_musculus' + assert variant.reference_name == genome.reference_name + assert variant.genome.sequence('1', 1, 4) == 'ACGT' + + +@pytest.mark.parametrize("damage,match", [ + (lambda resource: resource.update(path="../outside.gtf"), "inside"), + (lambda resource: resource.update(field="annotation_version"), "resource field"), + (lambda resource: resource.update(index=True), "scalar index"), + (lambda resource: resource.update(sha256="unknown"), "checksum"), +]) +def test_invalid_reference_manifest_rejected(tmp_path, reference, damage, match): + _, path = native_dataset(tmp_path, reference) + table = read_table(path) + payload = table.extra[DATASET_METADATA] + fragment = json.loads(payload["mutation_fragments"]["occurrence"]) + genome = fragment["variant"]["ensembl"] if "ensembl" in fragment["variant"] else fragment["variant"]["genome"] + damage(genome[REFERENCE_FIELD]["resources"][0]) + payload["mutation_fragments"]["occurrence"] = json.dumps(fragment) + with pytest.raises(ValueError, match=match): + NativeReferences(path).transform(payload, decode=True) + + +def test_native_cli_scores_and_constructs_replay_after_relocation(tmp_path, reference, monkeypatch): + import mhctools.cli + import pyensembl.download_cache + _, path = native_dataset(tmp_path, reference) + first = tmp_path / "first.tsv" + options = ["--vaccine-peptide-length", "10", "--no-processing-aware-annotation"] + run_cli(["--input-epitopes", str(path), "--output-epitopes", str(first)] + options) + # Move the original archive, then remove both raw annotations and indexes. + moved = tmp_path / "moved" + shutil.move(str(path.parent), str(moved)) + path = moved / "native.tsv" + delete_original_reference(tmp_path) + monkeypatch.setattr(mhctools.cli, "predictors_from_args", forbidden) + monkeypatch.setattr(pyensembl.download_cache.DownloadCache, "_fetch", forbidden) + output = tmp_path / "report" + run_cli(["--input-epitopes", str(path), "--index-native-references", + "--output-dir", str(output), "--output-epitopes", str(tmp_path / "second.tsv")] + options) + before, after = EpitopeDataset.load(first), EpitopeDataset.load(tmp_path / "second.tsv") + assert [e.per_allele_scores for e in before.epitopes] == [e.per_allele_scores for e in after.epitopes] + assert list(output.rglob("*.fasta")) and list(output.rglob("*.pdf")) + + +def test_index_option_requires_native_input(tmp_path): + from .test_epitope_dataset import evidence + table = tmp_path / "table.tsv" + evidence().to_tsv(table) + with pytest.raises(ValueError, match="requires a native"): + run_cli(["--input-topiary", str(table), "--index-native-references", + "--output-csv", str(tmp_path / "report.csv")]) diff --git a/vaxrank/cli/arg_parser.py b/vaxrank/cli/arg_parser.py index 83c541ce..8b3fdcbb 100644 --- a/vaxrank/cli/arg_parser.py +++ b/vaxrank/cli/arg_parser.py @@ -1023,6 +1023,10 @@ def _require_ensembl_release_for_template_reports(args): def add_external_rescoring_args(arg_parser): arg_parser.add_argument("--input-epitopes", default=None, metavar="FILE", help="Reload native epitope predictions and saved scoring evidence.") + arg_parser.add_argument( + "--index-native-references", action="store_true", default=False, + help="Prepare verified bundled annotation indexes for native input reports. " + "Uses local resources only; unnecessary for scoring or construct replay.") arg_parser.add_argument("--input-topiary", default=None, metavar="FILE", help="Load a normalized Topiary CSV/TSV using its original predictions.") arg_parser.add_argument( diff --git a/vaxrank/epitope_dataset.py b/vaxrank/epitope_dataset.py index db577ac9..fc52409e 100644 --- a/vaxrank/epitope_dataset.py +++ b/vaxrank/epitope_dataset.py @@ -55,6 +55,7 @@ class EpitopeDataset: construction_reports: list = field(default_factory=list) direct_sources: list = field(default_factory=list) policy_evaluations: list = field(default_factory=list) + native_references: object = field(default=None, repr=False, compare=False) def add_mutation(self, fragment, epitopes): """Capture direct predictions before vaccine windows discard context.""" @@ -254,6 +255,9 @@ def save(self, path): "direct_sources": self.direct_sources, "policy_evidence": encode_evaluations(self.policy_evaluations), } + from .native_references import NativeReferences + extra[DATASET_METADATA] = NativeReferences(path, self.native_references).transform( + extra[DATASET_METADATA]) result = TopiaryResult(self.result.df.copy(), metadata=self.result.metadata) result.extra = extra Path(path).parent.mkdir(parents=True, exist_ok=True) @@ -267,21 +271,41 @@ def load(cls, path): payload = result.extra.pop(DATASET_METADATA, None) if not isinstance(payload, dict) or payload.get("schema") != DATASET_SCHEMA: raise ValueError("Expected a native Vaxrank epitope dataset") + from .native_references import NativeReferences + references = NativeReferences(path) + payload = references.transform(payload, decode=True) epitopes = tuple(from_native_json(e, CandidateEpitope) for e in payload["epitopes"]) for epitope in epitopes: validate_peptide_alleles(epitope, str(path)) + antigens = {key: from_native_json(a, VaccineAntigen) + for key, a in payload["antigens"].items()} + fragments = {key: from_native_json(value, MutantProteinFragment) + for key, value in payload.get("mutation_fragments", {}).items()} + # Older datasets already persisted these IDs on the construction + # antigen. Reuse that recorded provenance only for the same transcripts; + # do not query an absent annotation to recreate it during replay. + for key, fragment in fragments.items(): + antigen = antigens.get(key) + transcript_ids = tuple(sorted({t.id for t in fragment.supporting_reference_transcripts})) + if (fragment.supporting_reference_protein_ids is None and antigen is not None + and antigen.transcript_ids == transcript_ids): + fragments[key] = replace(fragment, supporting_reference_protein_ids=antigen.protein_ids) return cls( result=result, epitopes=epitopes, provenance=tuple(from_native_json(p, InputProvenance) for p in payload["provenance"]), config=msgspec.convert(payload["config"], EpitopeConfig) if payload["config"] else None, - antigens={key: from_native_json(a, VaccineAntigen) - for key, a in payload["antigens"].items()}, + antigens=antigens, selection=payload["selection"], - mutation_fragments={key: from_native_json(value, MutantProteinFragment) - for key, value in payload.get("mutation_fragments", {}).items()}, + mutation_fragments=fragments, construction_reports=payload.get("construction_reports", []), direct_sources=payload.get("direct_sources", []), - policy_evaluations=decode_evaluations(payload.get("policy_evidence", []))) + policy_evaluations=decode_evaluations(payload.get("policy_evidence", [])), + native_references=references) + + def index_native_references(self): + """Prepare bundled local annotation indexes without acquiring data.""" + if self.native_references is not None: + self.native_references.index() def report_frame(self): """Report each occurrence/allele once; retain full evidence for scoring.""" diff --git a/vaxrank/external_rescoring.py b/vaxrank/external_rescoring.py index a13a1a18..1a25ccaa 100644 --- a/vaxrank/external_rescoring.py +++ b/vaxrank/external_rescoring.py @@ -482,6 +482,9 @@ def prepare_reports(inputs, epitope_config=None, *, mode='input', models=(), for key, antigen in r.dataset.antigens.items()} dataset.mutation_fragments = {key: fragment for r in updated if r.dataset is not None for key, fragment in r.dataset.mutation_fragments.items()} + from .native_references import NativeReferences + dataset.native_references = NativeReferences.combine( + r.dataset.native_references for r in updated if r.dataset is not None) from .native_serialization import to_native_json def native_row(row): frame = pd.DataFrame([row]).astype(object) @@ -545,6 +548,13 @@ def model_factory(): [(fmt, path) for fmt, path, _ in specs], scopes=[scope for _, _, scope in specs], manifest_path=getattr(args, 'input_manifest', None), genome=genome) + if getattr(args, 'index_native_references', False): + native = [report.dataset for report in loaded_reports + if report.source_format == 'epitopes' and report.dataset is not None] + if not native: + raise ValueError('--index-native-references requires a native epitope input') + for dataset in native: + dataset.index_native_references() from .config.loader import saved_construct_configuration, load_vaxrank_config saved_configs = [saved_construct_configuration(r.dataset.selection.get('run_configuration')) for r in loaded_reports if r.dataset is not None] diff --git a/vaxrank/mutant_protein_fragment.py b/vaxrank/mutant_protein_fragment.py index da633af6..ddefe6bb 100644 --- a/vaxrank/mutant_protein_fragment.py +++ b/vaxrank/mutant_protein_fragment.py @@ -328,6 +328,10 @@ class MutantProteinFragment(DataclassSerializable): # Blank for older archives and sources whose version was not supplied. sequence_source_version: str = "" + # Annotation provenance resolved while the source reference is available. + # None is unknown (legacy archives); () means no protein IDs were reported. + supporting_reference_protein_ids: Optional[tuple] = None + # Identity and provenance of the DNA outcome used to build this fragment. # Store data instead of Varcode's cyclic SV effect graph (varcode#438). # Defaults preserve compatibility with earlier serialized fragments. @@ -420,6 +424,9 @@ def from_isovar_result(cls, isovar_result): rna_evidence_subject=FRAGMENTS, sequence_source=SEQUENCE_SOURCE_ISOVAR_ASSEMBLY, sequence_source_version=isovar_version, + supporting_reference_protein_ids=tuple(sorted({ + t.protein_id for t in protein_sequence.transcripts + if getattr(t, "protein_id", None)})), supporting_reference_transcripts=protein_sequence.transcripts) @classmethod @@ -495,6 +502,8 @@ def from_variant_dna(cls, variant, protein_sequence_length): n_alt_reads_supporting_protein_sequence=0, sequence_source=SEQUENCE_SOURCE_VARCODE_TRANSLATION, supporting_reference_transcripts=transcripts, + supporting_reference_protein_ids=tuple(sorted({ + t.protein_id for t in transcripts if getattr(t, "protein_id", None)})), dna_effect_selection=selection, protein_start_offset=window_start, ) diff --git a/vaxrank/native_references.py b/vaxrank/native_references.py new file mode 100644 index 00000000..b9b63fde --- /dev/null +++ b/vaxrank/native_references.py @@ -0,0 +1,229 @@ +"""Verified relative resources for native custom annotation references.""" + +from copy import deepcopy +import hashlib +import json +from pathlib import Path +import shutil +import tempfile + +from pyensembl import Genome + +from .native_serialization import _json_loads + + +REFERENCE_FIELD = "__vaxrank_reference__" +REFERENCE_SCHEMA = "vaxrank.native_reference.v1" + + +def _digest(path): + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +class NativeReferences: + """Bundle available files; restore identities without acquiring resources.""" + + def __init__(self, path, previous=None): + self.root = Path(path).resolve().parent + self.directory = Path(Path(path).name + ".references") + self.availability = [] + self.genomes = {} + self._digests = {} + self._expected = {item['resolved_path']: item for item in + (previous.availability if previous is not None else ())} + + @classmethod + def combine(cls, references): + """Retain reference identities when native sources are composed.""" + references = [reference for reference in references if reference is not None] + if not references: + return None + combined = deepcopy(references[0]) + for reference in references[1:]: + combined.availability.extend(reference.availability) + combined.genomes.update(reference.genomes) + return combined + + def _file_digest(self, path): + stat = path.stat() + key = (path, stat.st_size, stat.st_mtime_ns, stat.st_ctime_ns) + if key not in self._digests: + self._digests[key] = _digest(path) + return self._digests[key] + + def _relative_path(self, relative): + path = Path(relative) + if path.is_absolute() or ".." in path.parts: + raise ValueError("Native reference path must stay inside the dataset directory") + resolved = (self.root / path).resolve() + if not resolved.is_relative_to(self.root): + raise ValueError("Native reference path escapes the dataset directory") + return resolved + + def _bundle_file(self, source): + digest = self._file_digest(source) + suffix = "".join(source.suffixes[-2:]) if source.suffix == ".gz" else source.suffix + relative = self.directory / (digest + suffix) + destination = self._relative_path(relative) + destination.parent.mkdir(parents=True, exist_ok=True) + if destination.exists(): + if self._file_digest(destination) != digest: + raise ValueError(f"Native reference bundle is corrupt: {destination}") + else: + with tempfile.NamedTemporaryFile(dir=destination.parent, delete=False) as stream: + temporary = Path(stream.name) + try: + shutil.copyfile(source, temporary) + if _digest(temporary) != digest: + raise ValueError(f"Reference changed during export: {source}") + temporary.replace(destination) + finally: + temporary.unlink(missing_ok=True) + return dict(path=str(relative), sha256=digest, size=destination.stat().st_size) + + def _resource(self, source): + expected = self._expected.get(str(source.resolve())) + if source.is_file(): + if expected is not None and (source.stat().st_size != expected['size'] + or self._file_digest(source) != expected['sha256']): + raise ValueError(f"Native reference checksum mismatch: {source}") + return self._bundle_file(source) + if expected is not None: + # Evidence-only exports retain the identity of absent resources; + # absence must not silently erase a previously recorded checksum. + return dict(path=str(self.directory / Path(expected['path']).name), + sha256=expected['sha256'], size=expected['size']) + return None + + def _encode_genome(self, value): + metadata = value["__class__"] + state = {key: item for key, item in value.items() if key != "__class__"} + resources = [] + if metadata["__name__"] == "Genome": + genome = Genome(**state) + # Public inspection does not download, copy, index or open a DB. + files = iter(genome.required_local_files()) + fields = [] + if state.get("genome_fasta_path_or_url"): + fields.append(("genome_fasta_path_or_url", None)) + if state.get("gtf_path_or_url"): + fields.append(("gtf_path_or_url", None)) + for key in ("transcript_fasta_paths_or_urls", "protein_fasta_paths_or_urls"): + fields.extend((key, index) for index, _ in enumerate(state.get(key) or [])) + for field, index in fields: + source = Path(next(files)) + resource = self._resource(source) + if resource is not None: + resources.append(dict(field=field, index=index, **resource)) + else: + # Exact Ensembl release/species is already portable. Attached local + # reference DNA is the only source path in its public state. + source = state.get("genome_fasta") + resource = self._resource(Path(source)) if isinstance(source, str) else None + if resource is not None: + resources.append(dict(field="genome_fasta", index=None, **resource)) + if resources: + value[REFERENCE_FIELD] = dict(schema=REFERENCE_SCHEMA, resources=resources) + return value + + def _decode_genome(self, value): + manifest = value.pop(REFERENCE_FIELD, None) + if manifest is None: + return value # Legacy references retain their exact identity. + if not isinstance(manifest, dict) or manifest.get("schema") != REFERENCE_SCHEMA: + raise ValueError("Unsupported native reference manifest") + resources = manifest.get("resources") + if not isinstance(resources, list) or not resources: + raise ValueError("Native reference manifest requires resources") + custom = value["__class__"]["__name__"] == "Genome" + seen = set() + for resource in resources: + if not isinstance(resource, dict) or set(resource) != {"field", "index", "path", "sha256", "size"}: + raise ValueError("Malformed native reference resource") + digest, size = resource["sha256"], resource["size"] + if (not isinstance(digest, str) or len(digest) != 64 + or any(c not in "0123456789abcdef" for c in digest) + or type(size) is not int or size < 0): + raise ValueError("Invalid native reference checksum or size") + field, index = resource["field"], resource["index"] + if field not in ({"gtf_path_or_url", "genome_fasta_path_or_url", + "transcript_fasta_paths_or_urls", "protein_fasta_paths_or_urls"} + if custom else {"genome_fasta"}): + raise ValueError("Invalid native reference resource field") + key = (field, index) + if key in seen: + raise ValueError("Duplicate native reference resource") + seen.add(key) + path = self._relative_path(resource["path"]) + available = path.is_file() + if available and (path.stat().st_size != resource["size"] + or self._file_digest(path) != resource["sha256"]): + raise ValueError(f"Native reference checksum mismatch: {path}") + self.availability.append(dict(resource, available=available, resolved_path=str(path))) + if field.endswith("_paths_or_urls"): + if type(index) is not int or not 0 <= index < len(value.get(field) or []): + raise ValueError("Invalid native reference resource index") + value[field][index] = str(path) + else: + if index is not None: + raise ValueError("Invalid native reference scalar index") + value[field] = str(path) + if custom: + # Derived indexes belong to this relocated bundle, never the old + # annotation cache. Their creation remains an explicit operation. + value["cache_directory_path"] = str(self.root / self.directory / "indexes") + value["copy_local_files_to_cache"] = False + value["decompress_on_download"] = False + state = {key: item for key, item in value.items() if key != "__class__"} + self.genomes[json.dumps(state, sort_keys=True)] = state + return value + + def _walk(self, value, operation): + if isinstance(value, list): + return [self._walk(item, operation) for item in value] + if not isinstance(value, dict): + return value + metadata = value.get("__class__", {}) + if isinstance(metadata, dict) and (metadata.get("__module__"), metadata.get("__name__")) in { + ("pyensembl.genome", "Genome"), ("pyensembl.ensembl_release", "EnsemblRelease")}: + value = operation(value) + return {key: self._walk(item, operation) for key, item in value.items()} + + def transform(self, payload, *, decode=False): + """Transform only native object fields, leaving source evidence intact.""" + operation = self._decode_genome if decode else self._encode_genome + def native(value): + return json.dumps(self._walk(_json_loads(value), operation)) if value is not None else None + payload = deepcopy(payload) + for field in ("epitopes", "provenance"): + payload[field] = [native(value) for value in payload[field]] + for field in ("antigens", "mutation_fragments"): + payload[field] = {key: native(value) for key, value in payload.get(field, {}).items()} + for source in payload.get("direct_sources", []): + source["variant"] = native(source["variant"]) + for report in payload.get("construction_reports", []): + report["genome"] = native(report.get("genome")) + report["rows"] = native(report["rows"]) + report["records"] = [native(value) for value in report["records"]] + return payload + + def index(self): + """Prepare verified local annotations explicitly, without downloads.""" + missing = [] + for item in self.availability: + path = Path(item["resolved_path"]) + if not path.is_file(): + missing.append(str(path)) + elif path.stat().st_size != item["size"] or self._file_digest(path) != item["sha256"]: + raise ValueError(f"Native reference checksum mismatch: {path}") + if missing: + raise ValueError("Native reference resources are unavailable: " + ", ".join(sorted(set(missing)))) + for state in self.genomes.values(): + genome = Genome(**state) + if not genome.required_local_files_exist(): + raise ValueError("Native annotation requires unavailable resources: " + str(genome)) + genome.index() diff --git a/vaxrank/vaccine_antigen.py b/vaxrank/vaccine_antigen.py index 01f78ca9..80126dc6 100644 --- a/vaxrank/vaccine_antigen.py +++ b/vaxrank/vaccine_antigen.py @@ -439,11 +439,13 @@ def from_mutant_protein_fragment(cls, fragment: Any) -> "VaccineAntigen": for transcript in transcripts if getattr(transcript, "transcript_id", "") })) - protein_ids = tuple(sorted({ - str(getattr(transcript, "protein_id", "")) - for transcript in transcripts - if getattr(transcript, "protein_id", "") - })) + protein_ids = getattr(fragment, "supporting_reference_protein_ids", None) + if protein_ids is None: + protein_ids = tuple(sorted({ + str(getattr(transcript, "protein_id", "")) + for transcript in transcripts + if getattr(transcript, "protein_id", "") + })) variant = getattr(fragment, "variant", None) source_identifier = str( getattr(variant, "short_description", "") or variant or "" diff --git a/vaxrank/version.py b/vaxrank/version.py index 6594f7f2..c9dd6604 100644 --- a/vaxrank/version.py +++ b/vaxrank/version.py @@ -1,4 +1,4 @@ -__version__ = "3.37.1" +__version__ = "3.37.2" def print_version(): From f2801fdbca00d4f12f50d7b8e44825b3a4d982ed Mon Sep 17 00:00:00 2001 From: Alex Rubinsteyn Date: Mon, 5 Oct 2026 14:39:27 -0400 Subject: [PATCH 2/2] Isolate attached DNA regression in a local release cache --- docs/native-references.md | 5 +++-- tests/test_native_references.py | 3 ++- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/docs/native-references.md b/docs/native-references.md index a8f401b4..2d14eb9f 100644 --- a/docs/native-references.md +++ b/docs/native-references.md @@ -38,8 +38,9 @@ vaxrank --input-epitopes result.tsv --index-native-references --output-dir repor ``` The equivalent Python call is `dataset.index_native_references()`. Indexing -verifies the resources again and acquires no data. Derived indexes are local to -the relocated bundle. Missing annotation resources produce an explicit error. +verifies the resources again and acquires no data. Custom-reference indexes are +local to the relocated bundle; standard Ensembl releases use PyEnsembl's normal +release cache. Missing annotation resources produce an explicit error. Earlier native files retain their original reference identities and now restore metadata lazily. Where the saved antigen already records protein IDs for the diff --git a/tests/test_native_references.py b/tests/test_native_references.py index 8aea6afd..e23e3904 100644 --- a/tests/test_native_references.py +++ b/tests/test_native_references.py @@ -177,11 +177,12 @@ def test_ensembl_release_preserves_identity_and_bundles_attached_dna(tmp_path, m from pyensembl import EnsemblRelease from varcode import Variant import pyensembl.download_cache + monkeypatch.setenv('PYENSEMBL_CACHE_DIR', str(tmp_path / 'release-cache')) monkeypatch.setattr(pyensembl.download_cache.DownloadCache, "_fetch", forbidden) dna = tmp_path / "dna.fa" dna.write_text(">1\nACGTACGTACGT\n") genome = EnsemblRelease(104, species="mouse", genome_fasta=str(dna)) - dataset = EpitopeDataset.from_predictions([_make_epitope("SIINFEKL")]) + dataset = EpitopeDataset.from_predictions([_make_epitope("SIINFEKL", ic50=100.)]) dataset.direct_sources = [{'variant': to_native_json(Variant('1', 2, 'C', 'T', ensembl=genome)), 'properties': {}}] path = tmp_path / "archive" / "native.tsv"