From fd67f844e3bb8a2d63f843dfd07f89c9acda0efc Mon Sep 17 00:00:00 2001 From: Alex Rubinsteyn Date: Thu, 1 Oct 2026 23:44:06 -0400 Subject: [PATCH 1/4] Score exact peptide occurrences with retained source context --- CHANGELOG.md | 11 ++ docs/api.md | 7 + docs/combined-sources.md | 65 +++++++++ tests/test_candidate_tables.py | 8 +- tests/test_consumer_workflows.py | 86 ++++++++++- tests/test_peptide_occurrences.py | 216 ++++++++++++++++++++++++++++ tests/test_twin_conformance.py | 14 ++ topiary/__init__.py | 4 +- topiary/candidates.py | 34 ++--- topiary/peptide_occurrences.py | 232 ++++++++++++++++++++++++++++++ topiary/predictor.py | 51 ++++++- 11 files changed, 695 insertions(+), 33 deletions(-) create mode 100644 tests/test_peptide_occurrences.py create mode 100644 topiary/peptide_occurrences.py diff --git a/CHANGELOG.md b/CHANGELOG.md index ed80dc6..d331056 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,17 @@ between published tags; older pre-5.0 notes are retained below. For current interfaces, see the [consumer guide](docs/consumer-guide.md). +## 5.91.0 + +- Score explicitly identified peptide occurrences with their original flanks, + coordinates and source evidence, without scanning additional windows. Share + identical inference inputs while preserving every observation; retain model, + allele-free/per-allele/haplotype scope and explicit comparator context (#367). +- Reuse occurrence prediction in additive table rescoring, batching compatible + peptides while retaining original measurements and selection behavior. +- Preserve the canonical prediction schema for empty named-peptide output, + including explicitly reported all-missing cache batches (#450). + ## 5.90.0 - Evaluate saved policies at explicit occurrence/allele identity before selecting diff --git a/docs/api.md b/docs/api.md index 13915ed..de88d36 100644 --- a/docs/api.md +++ b/docs/api.md @@ -91,11 +91,18 @@ return fresh read-only int8 arrays, so callers never share mutable matrix state. |--------|-------|----------| | `predict_from_named_sequences(dict)` | `{name: sequence}` | Sliding-window scan | | `predict_from_named_peptides(dict)` | `{name: peptide}` | Score as-is | +| `predict_from_peptide_occurrences(table)` | Unique `prediction_id`, peptide, optional flanks/coordinates/annotations | Score exact occurrences with their context; no scanning. | | `predict_from_sequences(list)` | `[sequence, ...]` | Sliding-window scan | | `predict_from_fragments(fragments)` | `[ProteinFragment]` | Universal path — any origin; fragment-level metadata and `target_intervals` threaded through. | | `predict_from_variants(variants)` | VariantCollection | Variant pipeline (builds `ProteinFragment`s internally and delegates). | | `predict_from_mutation_effects(effects)` | EffectCollection | Same as `predict_from_variants` but starting from pre-computed effects. | +`predict_peptide_occurrences(table, model, use_flanks=True, predict_wt=False)` +exposes the same occurrence prediction for one configured mhctools model. +It is exported from `topiary` and is also used by `rescore_candidates`. +See [exact occurrence prediction](combined-sources.md#predict-exact-peptide-occurrences) +for input identity, unknown flanks, comparator and coverage rules. + ## TopiaryResult `TopiaryResult` is the semantic result object for Topiary prediction tables. diff --git a/docs/combined-sources.md b/docs/combined-sources.md index 36cb472..e2ba2dd 100644 --- a/docs/combined-sources.md +++ b/docs/combined-sources.md @@ -534,6 +534,68 @@ after reloading. The composed workflow test covers original and additive scoring retained hypothesis selection, and count union without running a predictor during assembly. +## Predict exact peptide occurrences + +Topiary 5.91.0 adds `predict_peptide_occurrences` and +`TopiaryPredictor.predict_from_peptide_occurrences` for a selected peptide +universe. Supply one record per unique `prediction_id`, using the same occurrence +identity Vaxrank and the policy evaluator consume. Different occurrences of one +peptide remain separate even when they share a gene, coordinate or HLA candidate. + +```python +import pandas as pd +from topiary import ( + TopiaryPredictor, TopiaryResult, SelectionPolicy, evaluate_selection_policy, +) + +# model is an explicitly configured mhctools predictor with patient alleles. +occurrences = pd.DataFrame([ + dict(prediction_id="gene-a:12", peptide="SIINFEKL", peptide_offset=12, + n_flank="AAA", c_flank="GGG", gene="gene-a"), + dict(prediction_id="gene-b:20", peptide="SIINFEKL", peptide_offset=20, + n_flank="TTT", c_flank="CCC", gene="gene-b"), +]) +predictor = TopiaryPredictor(models=model) +predictions = predictor.predict_from_peptide_occurrences(occurrences) +evaluation = evaluate_selection_policy( + TopiaryResult(predictions), SelectionPolicy("binding", "1 / affinity.value"), + group_keys=["prediction_id", "peptide", "peptide_offset", "allele"], + kind_support=predictor.kind_support, +) +evaluation.evidence.to_tsv("occurrence-evidence.tsv") +``` + +The model scores only the supplied peptides against its configured alleles. +Within a shared sample/flank/genotype context, distinct peptides are batched and +identical inference inputs are scored once, then linked to each original +occurrence. Source annotations and read support are copied, never summed. +`prediction_mhc_dependence` records each kind's scope; haplotype predictions carry +the configured `allele_set`. An explicitly supplied haplotype set must match the +configured model. Missing kind/allele coverage and ambiguous outputs raise. + +Flank-dependent models require both `n_flank` and `c_flank`. An empty string is a +known terminus; a missing value is unknown. `use_flanks=False` explicitly requests +peptide-only inference while preserving the original context. The output's +`prediction_flanks_supplied` records whether flanks were supplied in the model call; it does not imply that every returned prediction kind depends on them. +Missing source offsets remain missing. Model peptide-length/domain checks still +apply; no padding, comparator sequence, or additional peptide window is inferred. + +With `TopiaryPredictor(predict_wt=True)` (or `predict_wt=True` on the standalone +function), supplied `wt_peptide` values are scored using their own +`wt_n_flank` / `wt_c_flank`. Flank-dependent comparator prediction requires those +flanks or explicit `use_flanks=False`; primary-peptide context is not borrowed. +Absent comparators keep missing scores. Explicit filter/sort settings operate on +occurrences; `only_novel_epitopes` does not reinterpret this already selected +peptide universe. Cache failures are strict by default; an explicit +`cache_miss_handler` / `on_miss` retains the existing reported-partial-batch +contract, with occurrence IDs in the failure records. + +Prediction inputs contain context and annotations, not historical measurement +columns. Keep imported measurements in their original table and use the additive +operation below. Context-aware live prediction requires mhctools' flank-aware +API; sparse historical tables with unknown model versions remain the separate +prediction-source work in [#368](https://github.com/openvax/topiary/issues/368). + ## Add prediction features on demand ```python @@ -564,6 +626,9 @@ models, versions, measurement semantics, UTC creation time, selection and discov Discovery source and prediction producer are separate facts. Configured mhctools models supply `predict_dataframe` and `kind_support()`. +This path reuses `predict_peptide_occurrences` to batch compatible requests and +preserve source identity; it appends primary-peptide features and leaves all +historical comparator measurements intact. Supplied flanks are forwarded by default; flank-dependent models require both flanks or explicit `use_flanks=False`. Haplotype models require a matching stated `allele_set`. Existing candidate identities are preserved; ORF-only rows diff --git a/tests/test_candidate_tables.py b/tests/test_candidate_tables.py index b73227f..1310cb7 100644 --- a/tests/test_candidate_tables.py +++ b/tests/test_candidate_tables.py @@ -123,7 +123,7 @@ def test_rescoring_only_selected_subset_preserves_original_ranking_and_adds_dsl_ model = Model() enriched = rescore_candidates(output, model, prefix="fresh", select="source_label == 'one'") feature = "fresh__testmodel__pMHC_affinity__value" - assert [call[0] for call in model.calls] == [["SIINFEKL"], ["GILGFVFTL"]] + assert [call[0] for call in model.calls] == [["SIINFEKL", "GILGFVFTL"]] assert enriched.df.loc[enriched.df.source_label.eq("two"), feature].isna().all() pd.testing.assert_frame_equal(enriched.df[before.columns], before) pd.testing.assert_frame_equal(output.df, before) @@ -149,7 +149,7 @@ def test_contexts_are_not_reused_and_flanks_reach_the_predictor(): output = combined(n_flank="OTHER") model = Model() scored = rescore_candidates(output, model, prefix="new") - assert len(model.calls) == 4 + assert len(model.calls) == 2 assert {c[1]["n_flanks"][0] for c in model.calls} == {"AAA", "OTHER"} assert scored.df.new__testmodel__pMHC_affinity__value.tolist() == [803, 13, 805, 15] @@ -158,7 +158,7 @@ def test_deduplicated_calls_keep_distinct_observations_and_evidence(): output = combined() model = Model() scored = rescore_candidates(output, model, prefix="new") - assert len(model.calls) == 2 + assert len(model.calls) == 1 assert len(scored.df) == len(output.df) assert scored.df.source_observation_id.nunique() == 4 @@ -327,7 +327,7 @@ def test_selective_rescoring_does_not_leak_across_genotypes(): feature = "new__testmodel__pMHC_affinity__value" assert scored.df[feature].iloc[:2].notna().all() assert scored.df[feature].iloc[2:].isna().all() - assert len(model.calls) == 2 + assert len(model.calls) == 1 def test_conflicting_predictions_need_source_or_run_identity_before_dsl_evaluation(): diff --git a/tests/test_consumer_workflows.py b/tests/test_consumer_workflows.py index 330cdf0..583e784 100644 --- a/tests/test_consumer_workflows.py +++ b/tests/test_consumer_workflows.py @@ -277,6 +277,85 @@ def test_peptide_level_measurement_conflicts_are_rejected_before_projection(): ) +@pytest.mark.parametrize("extension", ["csv", "tsv"]) +def test_explicit_occurrences_from_mixed_sources_rescore_and_replay(tmp_path, monkeypatch, extension): + from dataclasses import replace + from topiary import ( + ProteinFragment, SelectionPolicy, combine_sources, evaluate_selection_policy, + predict_peptide_occurrences, read_csv, read_tsv, rescore_candidates, replay_selection_policy, + select_policy_representatives, + ) + from .test_candidate_tables import Model + + fragments = [ + ProteinFragment(fragment_id="first", sequence="MAAASIINFEKLAAA", gene="G1", + reference_sequence="MAAAGIINFEKLAAA"), + ProteinFragment(fragment_id="second", sequence="MSIINFEKLGGG", gene="G2"), + ] + inputs = pd.DataFrame([dict( + prediction_id=fragment.fragment_id, fragment_id=fragment.fragment_id, + peptide=fragment.sequence[offset:offset + 8], peptide_offset=offset, + n_flank=fragment.sequence[max(0, offset - 3):offset], c_flank=fragment.sequence[offset + 8:offset + 11], + source_sequence=fragment.sequence, gene=fragment.gene, n_rna_alt=7, + wt_peptide=fragment.effective_baseline[offset:offset + 8] if fragment.effective_baseline else None, + wt_n_flank=fragment.effective_baseline[max(0, offset - 3):offset] if fragment.effective_baseline else None, + wt_c_flank=fragment.effective_baseline[offset + 8:offset + 11] if fragment.effective_baseline else None, + ) for fragment, offset in zip(fragments, [4, 1])]) + original = predict_peptide_occurrences(inputs, Model(), predict_wt=True) + assert original.peptide.tolist() == ["SIINFEKL", "SIINFEKL"] + assert original.wt_value.iloc[0] == 13. and pd.isna(original.wt_value.iloc[1]) + # Controlled original scores prefer the first occurrence; contextual fresh + # scores prefer the second. Neither operation changes the input peptides. + original["value"] = original["affinity"] = [10., 20.] + combined = combine_sources({"fragment": original, "lens": read_lens(LENS), + "pvacseq": read_pvacseq(PVACSEQ)}, sample_name="p") + before = combined.long_df.copy(deep=True) + policy = SelectionPolicy("original", "1 / affinity.value", filter_by="source_label == 'fragment'", duplicates="best") + def chosen(evaluation, field): + representatives = select_policy_representatives(evaluation) + return [evaluation.evidence.df.iloc[rows[0]][field] for rows in representatives.evidence_rows] + assert chosen(evaluate_selection_policy(combined, policy), "fragment_id") == ["first"] + + model = Model() + enriched = rescore_candidates(combined, model, prefix="fresh", select="source_label == 'fragment'") + assert [peptides for peptides, _ in model.calls] == [["SIINFEKL"], ["SIINFEKL"]] + pd.testing.assert_frame_equal(enriched.long_df[before.columns], before) + assert chosen(evaluate_selection_policy(enriched, policy), "fragment_id") == ["first"] + fresh = replace(policy, name="fresh", score_by="1 / fresh__testmodel__pMHC_affinity__value") + evaluated = evaluate_selection_policy(enriched, fresh) + second = select_policy_representatives(evaluated) + assert chosen(evaluated, "fragment_id") == ["second"] + assert chosen(evaluated, "n_rna_alt") == [7] + assert set(evaluated.evidence.df.source_label) == {"fragment", "lens", "pvacseq"} + assert evaluated.evidence.df.wt_peptide.dropna().str.len().gt(0).all() + + # Import-only records can also request exact peptide inference explicitly, + # without treating their unknown flanks as known termini. + imported = combined.df.loc[combined.df.source_label.ne("fragment")] + columns = [column for column in ("source_observation_id", "peptide", "peptide_offset", "n_flank", "c_flank", "wt_peptide") + if column in imported] + requests = imported[columns].drop_duplicates("source_observation_id").rename(columns={"source_observation_id": "prediction_id"}) + imported_model = Model() + with pytest.raises(ValueError, match="both flanks"): + predict_peptide_occurrences(requests, imported_model) + assert not imported_model.calls + exact = predict_peptide_occurrences(requests, imported_model, use_flanks=False) + assert set(exact.peptide) == set(requests.peptide) + assert exact.prediction_id.tolist() == requests.prediction_id.tolist() + assert sum(len(peptides) for peptides, _ in imported_model.calls) == requests.peptide.nunique() + pd.testing.assert_series_equal(exact.wt_peptide, requests.wt_peptide.reset_index(drop=True)) + + def forbidden(*args, **kwargs): + raise AssertionError("Saved evidence replay must not invoke inference") + monkeypatch.setattr(Model, "predict_dataframe", forbidden) + path = tmp_path / ("occurrence-policy." + extension) + getattr(evaluated.evidence, "to_" + extension)(path) + loaded = (read_csv if extension == "csv" else read_tsv)(path) + replay = replay_selection_policy(loaded) + pd.testing.assert_frame_equal(select_policy_representatives(replay), second, check_exact=True) + pd.testing.assert_frame_equal(replay.occurrences, evaluated.occurrences, check_exact=True) + + @pytest.mark.parametrize("wide", [False, True]) def test_source_tables_combine_rank_and_export_without_predicting(monkeypatch, tmp_path, wide, pandas_string_inference): from topiary import ( @@ -334,7 +413,8 @@ def forbidden(*args, **kwargs): model = Model() enriched = rescore_candidates(restored, model, prefix="fresh", select="source_label == 'direct'", use_flanks=False) - assert len(model.calls) == len(direct) + assert len(model.calls) == 1 + assert set(model.calls[0][0]) == set(direct.peptide) original_again = rank_candidates(enriched, expression, ascending=True, duplicates="best") assert original_again.candidate_id.tolist() == pooled.candidate_id.tolist() np.testing.assert_allclose(original_again.candidate_score, pooled.candidate_score, equal_nan=True) @@ -3281,7 +3361,7 @@ def forbidden(*args, **kwargs): result.df.iloc[2].rna_observations[0]])["value"] == 3 model = Model() scored = rescore_candidates(result, model, prefix="fresh") - assert len(model.calls) == 2 + assert len(model.calls) == 1 assert scored.df.candidate_id.nunique() == 2 saved = scored.to_wide() if wide else scored path = tmp_path / ("reconciled." + extension) @@ -3327,7 +3407,7 @@ def forbidden(*args, **kwargs): assert matches.iloc[0].candidate_id == ranked.iloc[0].candidate_id model = Model() rescored = rescore_candidates(combined, model, prefix="fresh", select="source_label == 'reported'") - assert len(model.calls) == 2 + assert len(model.calls) == 1 saved = rescored.to_wide() if wide else rescored path = tmp_path / ("exacto-combined." + extension) getattr(saved, "to_" + extension)(path) diff --git a/tests/test_peptide_occurrences.py b/tests/test_peptide_occurrences.py new file mode 100644 index 0000000..7c41ace --- /dev/null +++ b/tests/test_peptide_occurrences.py @@ -0,0 +1,216 @@ +"""Explicit occurrence inference preserves context without expanding peptides.""" +import random + +import numpy as np +import pandas as pd +import pytest +from mhctools import RandomBindingPredictor + +from topiary import ( + CachedPredictor, PartialPredictionWarning, SelectionPolicy, TopiaryPredictor, + TopiaryResult, combine_sources, evaluate_selection_policy, predict_peptide_occurrences, + rescore_candidates, read_tsv, +) +from .test_candidate_tables import Model, MultiKindModel, source +from .test_twin_conformance import PEPTIDE_OCCURRENCE_TWINS, NAMED_OCCURRENCE_TWINS + + +def occurrences(): + return pd.DataFrame([ + dict(prediction_id="one", peptide="SIINFEKL", n_flank="AAA", c_flank="GG", peptide_offset=3, gene="G1", n_rna_alt=7), + dict(prediction_id="other-gene", peptide="SIINFEKL", n_flank="AAA", c_flank="GG", peptide_offset=30, gene="G2", n_rna_alt=7), + dict(prediction_id="other-context", peptide="SIINFEKL", n_flank="T", c_flank="GG", peptide_offset=1, gene="G3", n_rna_alt=2), + dict(prediction_id="other-peptide", peptide="GILGFVFTL", n_flank="AAA", c_flank="GG", peptide_offset=3, gene="G4", n_rna_alt=3), + ]) + + +@pytest.mark.parametrize("predict", PEPTIDE_OCCURRENCE_TWINS) +def test_contexts_batch_once_and_keep_every_occurrence(predict): + inputs = occurrences() + before = inputs.copy(deep=True) + model = Model() + result = predict(inputs, model) + pd.testing.assert_frame_equal(inputs, before) + pd.testing.assert_frame_equal(result[inputs.columns], inputs) + assert result.value.tolist() == [803., 803., 801., 13.] + assert sum(len(peptides) for peptides, _ in model.calls) == 3 + assert len(model.calls) == 2 + assert result.prediction_flanks_supplied.all() + assert result.prediction_mhc_dependence.eq("single_allele").all() + # Batched and independent calls must have the same scores and evidence. + independent = pd.concat([predict(inputs.iloc[[index]], Model()) for index in range(len(inputs))], ignore_index=True) + pd.testing.assert_frame_equal(result, independent) + evaluated = evaluate_selection_policy(TopiaryResult(result), SelectionPolicy("context", "affinity.value"), + group_keys=["prediction_id", "peptide", "peptide_offset", "allele"]) + assert evaluated.occurrences.prediction_id.tolist() == inputs.prediction_id.tolist() + assert evaluated.occurrences.score.tolist() == result.value.tolist() + + +@pytest.mark.parametrize("peptides", [[], ["SIINFEKL"], ["SIINFEKL", "SIINFEKL", "GILGFVFTL"]]) +def test_real_mhctools_named_and_occurrence_doors_agree(peptides): + names = {str(index): peptide for index, peptide in enumerate(peptides)} + frames = [] + for predict in NAMED_OCCURRENCE_TWINS: + random.seed(123) + frames.append(predict(names, RandomBindingPredictor(alleles=["HLA-A*02:01", "HLA-B*07:02"]))) + columns = ["source_sequence_name", "peptide", "peptide_offset", "allele", "kind", "value", "score"] + left, right = (frame[columns].sort_values(columns[:4]).reset_index(drop=True) for frame in frames) + pd.testing.assert_frame_equal(left, right, check_dtype=False) + + +@pytest.mark.parametrize("change,match", [ + ({"prediction_id": "same"}, "unique"), + ({"prediction_id": None}, "prediction_id"), + ({"peptide": ""}, "peptide"), + ({"peptide_offset": -1}, "peptide_offset"), + ({"peptide_offset": 1.5}, "peptide_offset"), + ({"peptide_offset": True}, "peptide_offset"), + ({"peptide_length": 99}, "peptide_length"), + ({"value": 123.}, "prediction columns"), + ({"candidate_id": "old-candidate"}, "prediction columns"), + ({"n_flank": None}, "both flanks"), +]) +@pytest.mark.parametrize("predict", PEPTIDE_OCCURRENCE_TWINS) +def test_invalid_inputs_raise_before_inference(change, match, predict): + inputs = occurrences().assign(**change) + model = Model() + with pytest.raises(ValueError, match=match): + predict(inputs, model) + assert not model.calls + + +def test_unknown_and_empty_flanks_are_distinct_and_positions_not_invented(tmp_path): + inputs = pd.DataFrame([dict(prediction_id="known", peptide="SIINFEKL", n_flank="", c_flank=""), + dict(prediction_id="unknown", peptide="SIINFEKL")]) + output = predict_peptide_occurrences(inputs, Model(), use_flanks=False) + assert output.peptide_offset.isna().all() + assert output.n_flank.iloc[0] == "" and pd.isna(output.n_flank.iloc[1]) + assert not output.prediction_flanks_supplied.any() + path = tmp_path / "context.tsv" + TopiaryResult(output).to_tsv(path) + reloaded = read_tsv(path).df + assert reloaded.n_flank.iloc[0] == "" and pd.isna(reloaded.n_flank.iloc[1]) + assert reloaded.peptide_offset.isna().all() + with pytest.raises(ValueError, match="both flanks"): + predict_peptide_occurrences(inputs, Model()) + + +def test_model_scope_and_coverage_are_explicit(): + inputs = occurrences() + kinds = ("pMHC_affinity", "antigen_processing") + output = predict_peptide_occurrences(inputs, MultiKindModel(declared=kinds, emitted=kinds)) + assert len(output) == 8 + assert output.query("kind == 'antigen_processing'").allele.eq("").all() + assert output.query("kind == 'antigen_processing'").prediction_mhc_dependence.eq("none").all() + haplotype = Model(haplotype=True) + output = predict_peptide_occurrences(inputs, haplotype) + assert output.allele_set.eq("HLA-A*02:01").all() + with pytest.raises(ValueError, match="allele_set"): + predict_peptide_occurrences(inputs.assign(allele_set="HLA-B*07:02"), haplotype) + with pytest.raises(ValueError, match="pMHC_presentation"): + predict_peptide_occurrences(inputs, MultiKindModel(declared=(*kinds, "pMHC_presentation"))) + missing_allele = Model() + missing_allele.alleles = ["HLA-A*02:01", "HLA-B*07:02"] + with pytest.raises(ValueError, match="HLA-B\\*07:02"): + predict_peptide_occurrences(inputs, missing_allele) + with pytest.raises(ValueError, match="different peptide"): + predict_peptide_occurrences(inputs, Model(other=True)) + + +def test_filters_distinguish_occurrences_in_the_same_source(): + inputs = occurrences().assign(source_sequence_name="same-protein", peptide_offset=0) + predictor = TopiaryPredictor(models=Model(), filter_by="affinity.value < 802") + output = predictor.predict_from_peptide_occurrences(inputs) + assert output.prediction_id.tolist() == ["other-context", "other-peptide"] + + +@pytest.mark.parametrize("predict", PEPTIDE_OCCURRENCE_TWINS) +def test_empty_input_calls_no_model(predict): + model = Model() + assert predict([], model).empty + assert not model.calls + + +def test_explicit_comparators_use_only_their_own_context(): + inputs = occurrences().iloc[:2].copy() + inputs["wt_peptide"] = ["GILGFVFTL", None] + inputs["wt_n_flank"], inputs["wt_c_flank"] = "TTTTT", "" + model = Model() + output = TopiaryPredictor(models=model, predict_wt=True).predict_from_peptide_occurrences(inputs) + assert output.value.tolist() == [803., 803.] + assert output.wt_value.iloc[0] == 15. and pd.isna(output.wt_value.iloc[1]) + assert [peptides for peptides, _ in model.calls] == [["SIINFEKL"], ["GILGFVFTL"]] + assert model.calls[-1][1] == dict(n_flanks=["TTTTT"], c_flanks=[""]) + with pytest.raises(ValueError, match="both flanks"): + predict_peptide_occurrences(inputs.drop(columns="wt_n_flank"), Model(), predict_wt=True) + assert predict_peptide_occurrences(inputs.drop(columns="wt_n_flank"), Model(), predict_wt=True, use_flanks=False).wt_value.iloc[0] == 10. + + +def test_additive_rescoring_and_occurrence_predictions_agree(): + original = combine_sources({"one": source(), "two": source(n_flank="T")}, sample_name="p") + requests = original.df[["peptide", "n_flank", "c_flank", "candidate_sample"]].copy() + requests["prediction_id"] = original.df.source_observation_id + direct = predict_peptide_occurrences(requests, Model()) + enriched = rescore_candidates(original, Model(), prefix="fresh") + pd.testing.assert_series_equal(direct.value, enriched.df.fresh__testmodel__pMHC_affinity__value, check_names=False) + pd.testing.assert_frame_equal(enriched.df[original.df.columns], original.df) + + +def test_conflicting_outputs_cannot_be_relabelled_as_one_occurrence(): + class Conflicting(Model): + def predict_dataframe(self, peptides, **kwargs): + rows = super().predict_dataframe(peptides, **kwargs) + return pd.concat([rows, rows.assign(value=rows.value + 1)], ignore_index=True) + with pytest.raises(ValueError, match="Conflicting or ambiguous"): + predict_peptide_occurrences(occurrences(), Conflicting()) + + +def test_genotype_aliases_and_missing_values_use_mhcgnomes(): + model = Model(haplotype=True) + model.alleles = ["A0201"] + result = predict_peptide_occurrences(occurrences().assign(allele_set="HLA-A*02:01"), model) + assert result.allele_set.eq("HLA-A*02:01").all() + result = predict_peptide_occurrences(occurrences().assign(allele_set=pd.NA), model) + assert result.allele_set.eq("HLA-A*02:01").all() + with pytest.raises(TypeError, match="explicit booleans"): + predict_peptide_occurrences(occurrences(), model, use_flanks=None) + + +def test_partial_cache_coverage_is_reported_by_occurrence_and_empty_schema_survives(): + cache = CachedPredictor(source().df.iloc[:1].assign(n_flank="", c_flank="", peptide_length=8)) + inputs = occurrences() + reports = [] + predictor = TopiaryPredictor(models=cache, cache_miss_handler=reports.append) + with pytest.warns(PartialPredictionWarning, match="skipped 1"): + output = predictor.predict_from_peptide_occurrences(inputs, use_flanks=False) + assert output.prediction_id.tolist() == inputs.prediction_id.tolist()[:3] + assert reports[0]["source_sequence_name"] == "other-peptide" + assert reports[0]["stage"] == "occurrence" + assert reports[0]["model_key"] + assert reports[0]["configured_alleles"] == ["HLA-A*02:01"] + with pytest.warns(PartialPredictionWarning, match="skipped 1"): + empty = predictor.predict_from_peptide_occurrences(inputs.iloc[3:], use_flanks=False) + assert empty.empty and {"peptide", "prediction_id", "value", "kind", "allele"} <= set(empty) + + +def test_haplotype_comparator_can_have_a_different_deconvolved_presenter(): + class Haplotype(Model): + alleles = ["HLA-A*02:01", "HLA-B*07:02"] + def predict_dataframe(self, peptides, **kwargs): + output = super().predict_dataframe(peptides, **kwargs) + output["allele"] = np.where(output.peptide.eq("SIINFEKL"), self.alleles[0], self.alleles[1]) + return output + inputs = occurrences().iloc[:1].assign(wt_peptide="GILGFVFTL", wt_n_flank="T", wt_c_flank="") + output = predict_peptide_occurrences(inputs, Haplotype(haplotype=True), predict_wt=True) + assert output.allele.tolist() == ["HLA-A*02:01"] + assert output.allele_set.tolist() == ["HLA-A*02:01,HLA-B*07:02"] + assert output.wt_value.tolist() == [11.] + + +def test_shared_sequences_in_distinct_samples_remain_distinct_inference_contexts(): + inputs = occurrences().assign(sample_name=["one", "two", "one", "one"]) + model = Model() + output = predict_peptide_occurrences(inputs, model) + assert len(model.calls) == 3 + assert sum(len(peptides) for peptides, _ in model.calls) == 4 + pd.testing.assert_frame_equal(output[inputs.columns], inputs) diff --git a/tests/test_twin_conformance.py b/tests/test_twin_conformance.py index a5d7207..1c0ab30 100644 --- a/tests/test_twin_conformance.py +++ b/tests/test_twin_conformance.py @@ -50,6 +50,7 @@ combine_sources, rank_candidates, rank_with_policy, evaluate_scores, SelectionPolicy, resolve_selection_policy, evaluate_selection_policy, EvalContext, evaluate_filter, + predict_peptide_occurrences, Affinity, Column, apply_filter, apply_sort, ) from topiary.io_isovar import _check_isovar @@ -154,6 +155,19 @@ def _sv_support_record(support): OCCURRENCE_POLICY_TWINS = (evaluate_selection_policy, EvalContext) FILTER_DECISION_TWINS = (evaluate_filter, apply_filter) +# Explicit occurrence and named-peptide calls agree when context is the same. +# Driven by test_peptide_occurrences.py; source coordinates differ legitimately. +PEPTIDE_OCCURRENCE_TWINS = ( + predict_peptide_occurrences, + lambda occurrences, model, **kwargs: TopiaryPredictor(models=model).predict_from_peptide_occurrences(occurrences, **kwargs), +) +NAMED_OCCURRENCE_TWINS = ( + lambda names, model: TopiaryPredictor(models=model).predict_from_named_peptides(names), + lambda names, model: predict_peptide_occurrences( + pd.DataFrame([dict(prediction_id=name, peptide=peptide, peptide_offset=0) + for name, peptide in names.items()]), model, use_flanks=False), +) + # The CLI serializes the same exhaustive SV evidence policy as the public API. # Driven together in test_consumer_workflows, including absent protein rows. SV_INTEREST_REPORT_TWINS = (build_sv_interest_report, sv_interest_main) diff --git a/topiary/__init__.py b/topiary/__init__.py index 1b1abc7..0899510 100644 --- a/topiary/__init__.py +++ b/topiary/__init__.py @@ -17,6 +17,7 @@ reconcile_evidence, evidence_views, normalize_rna_observation, union_rna_observations, ) from .prediction_batch import predict_with_cache_miss_report, PartialPredictionWarning +from .peptide_occurrences import predict_peptide_occurrences from .predictor import ( TopiaryPredictor, fragment_from_effect, @@ -202,7 +203,7 @@ encode_amino_acids, ) -__version__ = "5.90.0" +__version__ = "5.91.0" __all__ = [ "normalize_isovar_rna_support", @@ -241,6 +242,7 @@ "write_selection_policy", "rank_with_policy", "rescore_candidates", + "predict_peptide_occurrences", "osteosarc_fixture_paths", "normalize_python_types", "TopiaryPredictor", diff --git a/topiary/candidates.py b/topiary/candidates.py index 3f68d58..9c740cc 100644 --- a/topiary/candidates.py +++ b/topiary/candidates.py @@ -13,7 +13,7 @@ from mhctools.pred import value_unit from .io_pvacseq import derive_mhc_class -from .predictor import from_predictions +from .peptide_occurrences import predict_peptide_occurrences from .ranking import ( apply_filter, as_dsl_node, evaluate_scores, format_allele_set, is_stated, mhc_dependence, split_allele_set, stated_values, @@ -417,27 +417,19 @@ def rescore_candidates(result, models, *, prefix, select=None, use_flanks=True): selected = selected[selected.candidate_id.notna()] selected = selected.drop_duplicates(["source_observation_id", "candidate_id", "allele_set"] if "allele_set" in selected else ["source_observation_id", "candidate_id"]) - features, descriptors, cache = {}, {}, {} - for _, row in selected.iterrows(): - n_flank, c_flank = row.get("n_flank"), row.get("c_flank") - flanks_known = isinstance(n_flank, str) and isinstance(c_flank, str) - genotype = format_allele_set(split_allele_set(row.get("allele_set"))) - for model_index, model in enumerate(models): - if use_flanks and getattr(model, "uses_flanking_sequences", False) and not flanks_known: - raise ValueError("Flank-dependent re-scoring requires both flanks; set use_flanks=False explicitly") - kwargs = dict(n_flanks=[n_flank], c_flanks=[c_flank]) if use_flanks and flanks_known else {} - cache_key = (model_index, row.candidate_sample, row.peptide, - n_flank if kwargs else None, c_flank if kwargs else None, genotype) - if cache_key not in cache: - raw = model.predict_dataframe([row.peptide], **kwargs) - support = model.kind_support() - predicted = from_predictions(raw) - if not predicted.empty and not predicted.peptide.eq(row.peptide).all(): - raise ValueError("Re-scoring returned a different peptide") - cache[cache_key] = (predicted, support) - predicted, support = cache[cache_key] + selected = selected.reset_index(drop=True) + requests = selected.reindex(columns=["peptide", "candidate_sample", "n_flank", "c_flank", "allele_set"]).copy() + requests["prediction_id"] = [str(index) for index in range(len(requests))] + features, descriptors = {}, {} + for model in models: + predicted = predict_peptide_occurrences(requests, model, use_flanks=use_flanks) + support = model.kind_support() + by_occurrence = {name: rows for name, rows in predicted.groupby("prediction_id", sort=False)} if not predicted.empty else {} + for index, row in selected.iterrows(): + genotype = format_allele_set(split_allele_set(row.get("allele_set"))) + predictions = by_occurrence[str(index)] matched_kinds = set() - for _, prediction in predicted.iterrows(): + for _, prediction in predictions.iterrows(): kind = str(prediction.kind) spec = next((v for k, v in support.items() if str(getattr(k, "value", k)) == kind), {}) dependence = spec.get("mhc_dependence") diff --git a/topiary/peptide_occurrences.py b/topiary/peptide_occurrences.py new file mode 100644 index 0000000..eddbc33 --- /dev/null +++ b/topiary/peptide_occurrences.py @@ -0,0 +1,232 @@ +"""Predict explicitly supplied peptide occurrences without scanning proteins.""" + +import numpy as np +import pandas as pd +import mhcgnomes + +from .predictor import from_predictions +from .prediction_batch import predict_with_cache_miss_report +from .ranking import format_allele_set, split_allele_set +from .wide import PREDICTION_COLUMNS, _concat_frames + + +def predict_peptide_occurrences(occurrences, model, *, use_flanks=True, predict_wt=False, on_miss=None): + """Score exact peptide occurrences with one configured mhctools model. + + Parameters + ---------- + occurrences : pandas.DataFrame or iterable of mappings + One row per unique, nonempty string ``prediction_id`` and ``peptide``. + Optional ``peptide_offset`` is a nonnegative source coordinate; absent + coordinates remain unknown. ``n_flank`` and ``c_flank`` are strings or + missing: an empty string states a known terminus. Other non-prediction + columns (source IDs, genes, RNA evidence, etc.) survive unchanged. + Prediction columns and allele-specific ``candidate_id`` / + ``candidate_allele`` / ``candidate_mhc_class`` are refused: keep historical measurements in their + original table and use :func:`rescore_candidates` for additive features. + A supplied ``allele_set`` must match a haplotype model's configuration. + model : mhctools predictor + Configured model exposing ``predict_dataframe`` and ``kind_support``. + Its alleles define the requested prediction scope. Every declared kind + must cover every input, including every configured allele for per-allele + kinds. Unsupported inputs, missing coverage and ambiguous outputs raise. + use_flanks : bool + Forward known flanks through mhctools' flank-aware API. Models declaring + ``uses_flanking_sequences`` require both flanks unless False explicitly + requests peptide-only inference. Original context is retained either way. + predict_wt : bool + Also score supplied ``wt_peptide`` comparators, using only explicitly + supplied ``wt_n_flank`` / ``wt_c_flank`` context. Missing comparators keep + missing scores. No baseline, coordinate correspondence or flank is + inferred from the primary peptide. + on_miss : callable, optional + Explicitly report and skip cache coverage failures, using + :func:`predict_with_cache_miss_report`. The report's source name is the + occurrence ID. Other errors still raise; partial results emit a warning. + + Returns + ------- + pandas.DataFrame + Topiary long-form predictions in occurrence order, with model/version, + ``prediction_mhc_dependence``, ``prediction_flanks_supplied`` and genotype + scope. Equal inputs within a sample/flank/genotype partition are inferred + once and expanded back to all occurrences; biological support is never + summed. Empty input calls no model and returns an empty prediction frame. + No sliding windows, additional peptide sequences or source candidates + are created. Model errors propagate; input tables are never mutated. + """ + if not isinstance(use_flanks, bool) or not isinstance(predict_wt, bool): + raise TypeError("use_flanks and predict_wt must be explicit booleans") + if on_miss is not None and not callable(on_miss): + raise TypeError("on_miss must be callable or None") + frame = pd.DataFrame(occurrences).copy().reset_index(drop=True) + if not frame.columns.is_unique: + raise ValueError("Occurrence columns must be unique") + if frame.empty: + return from_predictions([], extra_columns={column: [] for column in frame}) + for column in ("prediction_id", "peptide"): + if column not in frame or not frame[column].map(lambda value: isinstance(value, str) and bool(value.strip())).all(): + raise ValueError(f"Occurrences require nonempty string {column}") + if frame.prediction_id.duplicated().any(): + raise ValueError("prediction_id must be unique; retain distinct source observations under distinct IDs") + reserved = (PREDICTION_COLUMNS | {"allele", "value_unit", "measurement_context", "predictor_name", + "offset", "prediction_mhc_dependence", "prediction_flanks_supplied", + "candidate_id", "candidate_allele", "candidate_mhc_class"}) & set(frame) + if reserved: + raise ValueError(f"Occurrence input contains prediction columns: {sorted(reserved)}") + if not callable(getattr(model, "predict_dataframe", None)) or not callable(getattr(model, "kind_support", None)): + raise TypeError("model must expose predict_dataframe and kind_support") + support = {str(getattr(kind, "value", kind)): spec for kind, spec in model.kind_support().items()} + if not support or any(spec.get("mhc_dependence") not in {"none", "single_allele", "haplotype"} + for spec in support.values()): + raise ValueError("Model must declare MHC dependence for every prediction kind") + alleles = sorted({mhcgnomes.parse(str(allele)).to_string() + for allele in (getattr(model, "alleles", ()) or ())}) + if any(spec["mhc_dependence"] != "none" for spec in support.values()) and not alleles: + raise ValueError("Allele-dependent models require configured alleles") + for column in ("n_flank", "c_flank"): + if column not in frame: + frame[column] = None + if not frame[column].map(lambda value: isinstance(value, str) or + (pd.api.types.is_scalar(value) and pd.isna(value))).all(): + raise ValueError(f"{column} must contain strings or missing values") + if "peptide_offset" not in frame: + frame["peptide_offset"] = pd.array([None] * len(frame), dtype="Int64") + offsets = frame.peptide_offset.dropna() + if not offsets.map(lambda value: not isinstance(value, (bool, np.bool_)) + and isinstance(value, (int, float, np.integer, np.floating)) + and np.isfinite(value) and value >= 0 and value == int(value)).all(): + raise ValueError("peptide_offset must be a nonnegative integer or missing") + if "peptide_length" in frame and not frame.peptide_length.eq(frame.peptide.str.len()).all(): + raise ValueError("peptide_length disagrees with peptide") + flank_model = bool(getattr(model, "uses_flanking_sequences", False)) + has_flanks = frame.n_flank.map(lambda value: isinstance(value, str)) & frame.c_flank.map(lambda value: isinstance(value, str)) + if use_flanks and flank_model and not has_flanks.all(): + raise ValueError("Flank-dependent prediction requires both flanks; set use_flanks=False explicitly") + if predict_wt and "wt_peptide" in frame: + comparators = frame.loc[frame.wt_peptide.notna()] + if not comparators.wt_peptide.map(lambda value: isinstance(value, str) and bool(value.strip())).all(): + raise ValueError("wt_peptide must be a nonempty string or missing") + if use_flanks and flank_model and not comparators.empty: + if any(column not in comparators or not comparators[column].map(lambda value: isinstance(value, str)).all() + for column in ("wt_n_flank", "wt_c_flank")): + raise ValueError("Flank-dependent comparator prediction requires both flanks for the comparator") + genotypes = frame.get("allele_set", pd.Series(None, index=frame.index, dtype=object)).map( + lambda value: "" if pd.isna(value) else format_allele_set( + mhcgnomes.parse(name).to_string() for name in split_allele_set(value))) + if any(spec["mhc_dependence"] == "haplotype" for spec in support.values()): + if any(genotype and genotype != format_allele_set(alleles) for genotype in genotypes): + raise ValueError("Haplotype occurrence allele_set must match the configured model") + if on_miss is not None: + reported = predict_with_cache_miss_report( + lambda inputs: predict_peptide_occurrences( + frame.loc[frame.prediction_id.isin(inputs)], model, + use_flanks=use_flanks, predict_wt=predict_wt), + dict(zip(frame.prediction_id, frame.peptide)), on_miss=on_miss, + context=dict(stage="occurrence", prediction_method_name=str( + getattr(model, "prediction_method_name", getattr(model, "predictor_name", type(model).__name__))), + predictor_version=str(getattr(model, "predictor_version", "")), + configured_alleles=alleles, use_flanks=use_flanks)) + return reported if not reported.empty else predict_peptide_occurrences(frame.iloc[:0], model) + + partitions = {} + for index, row in frame.iterrows(): + flanks = (row.n_flank, row.c_flank) if use_flanks and has_flanks.iloc[index] else (None, None) + samples = tuple(None if pd.isna(row.get(column)) else row[column] + for column in ("sample_name", "candidate_sample")) + partitions.setdefault((*samples, *flanks, genotypes.iloc[index]), []).append(index) + expanded = [] + for (*_, n_flank, c_flank, genotype), indices in partitions.items(): + inputs = frame.iloc[indices] + peptides = inputs.peptide.drop_duplicates().tolist() + kwargs = {} if n_flank is None else dict(n_flanks=[n_flank] * len(peptides), c_flanks=[c_flank] * len(peptides)) + predicted = from_predictions(model.predict_dataframe(peptides, **kwargs)) + if predicted.empty: + raise ValueError(f"Model returned no prediction for {peptides}") + if not predicted.peptide.isin(peptides).all(): + raise ValueError("Model returned a different peptide than the requested occurrences") + if not predicted.kind.isin(support).all(): + raise ValueError("Model returned an undeclared prediction kind") + predicted["allele"] = predicted.allele.map( + lambda value: mhcgnomes.parse(str(value)).to_string() if isinstance(value, str) and value.strip() else "") + for peptide in peptides: + for kind, spec in support.items(): + rows = predicted[predicted.peptide.eq(peptide) & predicted.kind.eq(kind)] + dependence = spec["mhc_dependence"] + expected = set(alleles) if dependence == "single_allele" else None + actual = set(rows.allele) + if rows.empty or (expected is not None and expected != actual): + raise ValueError(f"Model returned no complete {kind} prediction for {peptide}/{','.join(alleles)}") + if dependence == "none" and actual != {""}: + raise ValueError(f"Allele-free {kind} output must not name an allele") + if dependence == "haplotype" and not actual <= set(alleles) | {""}: + raise ValueError("Haplotype output names an allele outside the configured genotype") + if rows.duplicated(["allele"]).any() or (dependence != "single_allele" and len(rows) != 1): + raise ValueError(f"Conflicting or ambiguous {kind} predictions for {peptide}") + for column, context in (("n_flank", n_flank), ("c_flank", c_flank)): + if column in predicted: + reported = predicted[column].dropna() + if context is None: + matches = reported.eq("") + elif flank_model: + matches = reported.map(lambda value: isinstance(value, str) and + (context.endswith(value) if column == "n_flank" else context.startswith(value))) + else: + matches = pd.Series(True, index=reported.index) + if not matches.all(): + raise ValueError("Returned flank context does not match the requested occurrence") + predicted["prediction_mhc_dependence"] = predicted.kind.map(lambda kind: support[kind]["mhc_dependence"]) + predicted["prediction_flanks_supplied"] = bool(kwargs) + if "allele_set" in predicted: + reported = predicted.loc[predicted.prediction_mhc_dependence.eq("haplotype"), "allele_set"] + if any(value and value != format_allele_set(alleles) + for value in reported.map(lambda value: "" if pd.isna(value) else format_allele_set( + mhcgnomes.parse(name).to_string() for name in split_allele_set(value)))): + raise ValueError("Returned haplotype allele_set differs from the configured model") + predicted["allele_set"] = np.where(predicted.prediction_mhc_dependence.eq("haplotype"), + format_allele_set(alleles), genotype) + # Input context is authoritative; prediction values remain model-owned. + payload = predicted.drop(columns=[column for column in frame if column != "peptide"] + + ["source_sequence_name", "peptide_offset", "sample_name"], errors="ignore") + if "allele_set" in frame: + payload["allele_set"] = predicted.allele_set + context = inputs.drop(columns=["allele_set"], errors="ignore") + expanded.append(context.merge(payload, on="peptide", how="left", validate="many_to_many")) + output = _concat_frames(expanded) + order = {name: i for i, name in enumerate(frame.prediction_id)} + output = output.iloc[np.argsort(output.prediction_id.map(order), kind="stable")].reset_index(drop=True) + if "source_sequence_name" not in output: + output["source_sequence_name"] = output.prediction_id + if "sample_name" not in output: + output["sample_name"] = "" + if predict_wt: + output = _attach_comparators(output, frame, model, use_flanks) + return output + + +def _attach_comparators(output, occurrences, model, use_flanks): + """Plumbing for the explicit comparator pass of predict_peptide_occurrences.""" + present = occurrences.get("wt_peptide", pd.Series(None, index=occurrences.index, dtype=object)).map( + lambda value: isinstance(value, str) and bool(value)) + fields = ("value", "score", "affinity", "percentile_rank", "prediction_method_name", "predictor_version") + if not present.any(): + for field in fields: + output["wt_" + field] = np.nan + return output + comparators = occurrences.loc[present].copy() + comparators["peptide"] = comparators.wt_peptide + comparators = comparators.drop(columns=["peptide_length"], errors="ignore") + for column in ("n_flank", "c_flank", "peptide_offset"): + comparators[column] = comparators.get("wt_" + column) + comparators = comparators.drop(columns=[column for column in comparators if column.startswith("wt_")]) + scores = predict_peptide_occurrences(comparators, model, use_flanks=use_flanks) + # A haplotype's deconvolved presenter may change for the comparator. Its + # score still describes the same configured genotype, not that one allele. + def key(row): + return (row.prediction_id, row.kind, + row.allele if row.prediction_mhc_dependence == "single_allele" else "") + by_scope = {key(row): row for row in scores.itertuples(index=False)} + matches = [by_scope.get(key(row)) for row in output.itertuples(index=False)] + for field in fields: + output["wt_" + field] = [getattr(row, field, np.nan) for row in matches] + return output diff --git a/topiary/predictor.py b/topiary/predictor.py index 17cab08..e455cf5 100644 --- a/topiary/predictor.py +++ b/topiary/predictor.py @@ -1005,6 +1005,43 @@ def predict_from_named_peptides(self, name_to_peptide_dict): self._strip_internal_columns(self._apply_filter(df)) ) + def predict_from_peptide_occurrences(self, occurrences, *, use_flanks=True): + """Predict exact occurrences with their source context, without scanning. + + Parameters + ---------- + occurrences : pandas.DataFrame or iterable of mappings + Explicit unique ``prediction_id`` and peptide, with optional source + coordinates, flanks and non-prediction annotations. See + :func:`predict_peptide_occurrences` for validation and scope rules. + use_flanks : bool + Use supplied flanks. False explicitly requests peptide-only scores, + while preserving original occurrence context in the output. + + Returns + ------- + pandas.DataFrame + Configured-model predictions with all occurrence identities. Empty + input calls no predictor. Explicit filter/sort settings operate at + occurrence scope. As for named peptides, ``only_novel_epitopes`` + does not reinterpret the caller's already-selected peptide universe. + ``predict_wt`` scores only supplied comparators with their supplied + context; ``cache_miss_handler`` explicitly enables partial results. + """ + from .peptide_occurrences import predict_peptide_occurrences + occurrences = pd.DataFrame(occurrences) + frames = [] + for model, model_key in zip(self.models, self._model_keys): + handler = None if self.cache_miss_handler is None else ( + lambda report, key=model_key: self.cache_miss_handler(dict(report, model_key=key))) + frames.append(self._attach_model_key(predict_peptide_occurrences( + occurrences, model, use_flanks=use_flanks, predict_wt=self.predict_wt, + on_miss=handler), model_key)) + df = _concat_frames(frames) if frames else pd.DataFrame() + keys = ["prediction_id", "peptide", "peptide_offset", "allele"] + return self._attach_result_attrs(self._strip_internal_columns( + self._apply_filter(df, group_keys=keys)).reset_index(drop=True)) + def _predict_raw(self, name_to_sequence_dict): """Run models and format output, without applying filter/ranking.""" for model in self.models: @@ -1059,7 +1096,13 @@ def _predict_raw_peptides_for_model( context=self._cache_miss_context(model, model_key, stage)) # All inputs may have been reported as missing. Preserve an empty # table schema just as the protein path does, including for CSV output. - return self._format_prediction_df(result) if result.empty else result + if result.empty: + empty = from_predictions([]) + for column in result: + if column not in empty: + empty[column] = result[column] + return empty + return result def _cache_miss_context(self, model, model_key, stage): """Provenance for one configured model's prediction pass.""" @@ -1206,7 +1249,7 @@ def _format_prediction_df(self, df): """Normalize mhctools prediction output to Topiary's schema.""" return _normalize_prediction_frame(df) - def _apply_filter(self, df): + def _apply_filter(self, df, *, group_keys=None): """Apply filter and sort if configured.""" if df.empty: return df @@ -1226,12 +1269,12 @@ def _apply_filter(self, df): if self.filter_by is not None: df = apply_filter( df, self.filter_by, kind_support=kind_support, - default_methods=default_methods, + default_methods=default_methods, group_keys=group_keys, ) if self.sort_by: df = apply_sort( df, self.sort_by, sort_direction=self.sort_direction, - kind_support=kind_support, default_methods=default_methods, + kind_support=kind_support, default_methods=default_methods, group_keys=group_keys, ) return df From 0d427b7a51081bf442e44899ba9f28d38532d3f2 Mon Sep 17 00:00:00 2001 From: Alex Rubinsteyn Date: Thu, 1 Oct 2026 23:53:01 -0400 Subject: [PATCH 2/4] Match both wildtype prediction paths by declared MHC scope --- CHANGELOG.md | 3 ++ docs/api.md | 6 +++ tests/test_peptide_occurrences.py | 62 +++++++++++++++++++++++++++++-- tests/test_twin_conformance.py | 9 +++++ topiary/__init__.py | 2 + topiary/peptide_occurrences.py | 6 +-- topiary/predictor.py | 23 ++++++++++-- topiary/ranking/__init__.py | 2 + topiary/ranking/nodes.py | 41 ++++++++++++++++++++ 9 files changed, 144 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d331056..3919777 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,9 @@ For current interfaces, see the [consumer guide](docs/consumer-guide.md). peptides while retaining original measurements and selection behavior. - Preserve the canonical prediction schema for empty named-peptide output, including explicitly reported all-missing cache batches (#450). +- Match fragment and explicit-occurrence WT scores by declared MHC scope, so + a different haplotype presenter does not erase the comparator score. Export + `prediction_mhc_scope` for consumers sharing that identity rule (#453). ## 5.90.0 diff --git a/docs/api.md b/docs/api.md index de88d36..f9ad333 100644 --- a/docs/api.md +++ b/docs/api.md @@ -103,6 +103,12 @@ It is exported from `topiary` and is also used by `rescore_candidates`. See [exact occurrence prediction](combined-sources.md#predict-exact-peptide-occurrences) for input identity, unknown flanks, comparator and coverage rules. +`prediction_mhc_scope(allele, dependence=..., allele_set=...)` returns a canonical +MHC identity tuple: one allele, the full genotype, or an allele-free scope. +Missing required identity returns `None`, which must remain unmatchable. +Comparator joins use this public rule in both fragment and occurrence prediction; +a haplotype's selected presenter is not its genotype identity. + ## TopiaryResult `TopiaryResult` is the semantic result object for Topiary prediction tables. diff --git a/tests/test_peptide_occurrences.py b/tests/test_peptide_occurrences.py index 7c41ace..99400a5 100644 --- a/tests/test_peptide_occurrences.py +++ b/tests/test_peptide_occurrences.py @@ -7,12 +7,12 @@ from mhctools import RandomBindingPredictor from topiary import ( - CachedPredictor, PartialPredictionWarning, SelectionPolicy, TopiaryPredictor, + CachedPredictor, PartialPredictionWarning, ProteinFragment, SelectionPolicy, TopiaryPredictor, TopiaryResult, combine_sources, evaluate_selection_policy, predict_peptide_occurrences, - rescore_candidates, read_tsv, + rescore_candidates, read_tsv, prediction_mhc_scope, ) from .test_candidate_tables import Model, MultiKindModel, source -from .test_twin_conformance import PEPTIDE_OCCURRENCE_TWINS, NAMED_OCCURRENCE_TWINS +from .test_twin_conformance import PEPTIDE_OCCURRENCE_TWINS, NAMED_OCCURRENCE_TWINS, WILDTYPE_SCOPE_TWINS def occurrences(): @@ -214,3 +214,59 @@ def test_shared_sequences_in_distinct_samples_remain_distinct_inference_contexts assert len(model.calls) == 3 assert sum(len(peptides) for peptides, _ in model.calls) == 4 pd.testing.assert_frame_equal(output[inputs.columns], inputs) + + +@pytest.mark.parametrize("dependence", ["single_allele", "haplotype", "none"]) +def test_fragment_and_exact_comparators_match_the_same_mhc_scope(dependence): + class ScopeModel: + alleles = ["HLA-A*02:01", "HLA-B*07:02"] + default_peptide_lengths = [8] + uses_flanking_sequences = True + + def kind_support(self): + return {"pMHC_presentation" if dependence != "none" else "antigen_processing": { + "mhc_dependence": dependence}} + + def predict_dataframe(self, peptides, **kwargs): + kind = next(iter(self.kind_support())) + records = [] + for peptide in peptides: + mutant = peptide.startswith("S") + alleles = self.alleles if dependence == "single_allele" else ( + [self.alleles[0 if mutant else 1]] if dependence == "haplotype" else [None]) + for allele in alleles: + records.append(dict(peptide=peptide, allele=allele, kind=kind, + score=.2 if mutant else .8, value=.2 if mutant else .8, + percentile_rank=None, n_flank="", c_flank="", + prediction_method_name="scope_fixture", predictor_version="1")) + return pd.DataFrame(records) + + def predict_proteins_dataframe(self, inputs): + assert all(len(sequence) == 8 for sequence in inputs.values()) + return pd.concat([self.predict_dataframe([sequence]).assign(source_sequence_name=name, peptide_offset=0) + for name, sequence in inputs.items()], ignore_index=True) + + fragment = ProteinFragment(fragment_id="one", sequence="SIINFEKL", reference_sequence="GIINFEKL") + frames = [predict(fragment, ScopeModel()) for predict in WILDTYPE_SCOPE_TWINS] + for output in frames: + assert output.wt_value.eq(.8).all() + assert output.value.eq(.2).all() + if dependence == "haplotype": + assert output.allele.tolist() == ["HLA-A*02:01"] + assert output.allele_set.tolist() == ["HLA-A*02:01,HLA-B*07:02"] + columns = ["peptide", "kind", "value", "wt_value", "wt_score", "wt_prediction_method_name", "wt_predictor_version"] + pd.testing.assert_frame_equal(frames[0][columns], frames[1][columns]) + + +def test_public_mhc_scope_normalizes_identity_and_does_not_match_unknowns(): + assert prediction_mhc_scope("A0201", dependence="single_allele") == prediction_mhc_scope("HLA-A*02:01", dependence="single_allele") + assert prediction_mhc_scope("H2-Kb", dependence="single_allele") == prediction_mhc_scope("H-2-Kb", dependence="single_allele") + first = prediction_mhc_scope("A0201", dependence="haplotype", allele_set="A0201,B0702") + same = prediction_mhc_scope("B0702", dependence="haplotype", allele_set=["B0702", "A0201"]) + assert first == same + assert first != prediction_mhc_scope("A0201", dependence="haplotype", allele_set=["A0201"]) + assert prediction_mhc_scope(None, dependence="single_allele") is None + assert prediction_mhc_scope("A0201", dependence="haplotype") is None + assert prediction_mhc_scope(None, dependence="none") == ("none",) + with pytest.raises(ValueError, match="Unknown MHC dependence"): + prediction_mhc_scope("A0201", dependence="unknown") diff --git a/tests/test_twin_conformance.py b/tests/test_twin_conformance.py index 1c0ab30..6a5b77e 100644 --- a/tests/test_twin_conformance.py +++ b/tests/test_twin_conformance.py @@ -168,6 +168,15 @@ def _sv_support_record(support): for name, peptide in names.items()]), model, use_flanks=False), ) +# Both comparator paths must match measurements by their declared MHC scope, +# including haplotypes whose MT/WT deconvolved presenter differs (#453). +WILDTYPE_SCOPE_TWINS = ( + lambda fragment, model: TopiaryPredictor(models=model, predict_wt=True, only_novel_epitopes=False).predict_from_fragments([fragment]), + lambda fragment, model: TopiaryPredictor(models=model, predict_wt=True).predict_from_peptide_occurrences([ + dict(prediction_id=fragment.fragment_id, peptide=fragment.sequence, peptide_offset=0, + n_flank="", c_flank="", wt_peptide=fragment.effective_baseline, wt_n_flank="", wt_c_flank="")]), +) + # The CLI serializes the same exhaustive SV evidence policy as the public API. # Driven together in test_consumer_workflows, including absent protein rows. SV_INTEREST_REPORT_TWINS = (build_sv_interest_report, sv_interest_main) diff --git a/topiary/__init__.py b/topiary/__init__.py index 0899510..4f41bd9 100644 --- a/topiary/__init__.py +++ b/topiary/__init__.py @@ -44,6 +44,7 @@ KIND_MHC_DEPENDENCE, MHC_DEPENDENCE_VALUES, mhc_dependence, + prediction_mhc_scope, resolve_kind, resolve_default_methods, describe_default_versions, @@ -243,6 +244,7 @@ "rank_with_policy", "rescore_candidates", "predict_peptide_occurrences", + "prediction_mhc_scope", "osteosarc_fixture_paths", "normalize_python_types", "TopiaryPredictor", diff --git a/topiary/peptide_occurrences.py b/topiary/peptide_occurrences.py index eddbc33..74928e8 100644 --- a/topiary/peptide_occurrences.py +++ b/topiary/peptide_occurrences.py @@ -6,7 +6,7 @@ from .predictor import from_predictions from .prediction_batch import predict_with_cache_miss_report -from .ranking import format_allele_set, split_allele_set +from .ranking import format_allele_set, split_allele_set, prediction_mhc_scope from .wide import PREDICTION_COLUMNS, _concat_frames @@ -223,8 +223,8 @@ def _attach_comparators(output, occurrences, model, use_flanks): # A haplotype's deconvolved presenter may change for the comparator. Its # score still describes the same configured genotype, not that one allele. def key(row): - return (row.prediction_id, row.kind, - row.allele if row.prediction_mhc_dependence == "single_allele" else "") + return (row.prediction_id, row.kind, prediction_mhc_scope( + row.allele, dependence=row.prediction_mhc_dependence, allele_set=row.allele_set)) by_scope = {key(row): row for row in scores.itertuples(index=False)} matches = [by_scope.get(key(row)) for row in output.itertuples(index=False)] for field in fields: diff --git a/topiary/predictor.py b/topiary/predictor.py index e455cf5..edf7ce2 100644 --- a/topiary/predictor.py +++ b/topiary/predictor.py @@ -32,6 +32,8 @@ apply_sort, as_dsl_node, as_dsl_nodes, + mhc_dependence, + prediction_mhc_scope, ) from .io import _model_version_str from .wide import _concat_frames @@ -1376,7 +1378,8 @@ def _maybe_predict_wt_peptides(self, df, fragments=None): ``wt_peptide`` itself is derived while building fragment rows. This helper performs the optional second prediction pass and joins each WT prediction back to the matching mutant row by - allele, peptide length, prediction kind, and predictor identity. + declared MHC scope, peptide length, prediction kind and predictor + identity. A haplotype's deconvolved presenter is not its scope. """ if not self.predict_wt or df.empty or "wt_peptide" not in df.columns: return df @@ -1411,7 +1414,6 @@ def _ensure_wt_columns(out): valid = ( df["wt_peptide"].notna() - & df["allele"].notna() & df["wt_peptide_length"].notna() & df["fragment_id"].notna() & df[_WT_OFFSET_COLUMN].notna() @@ -1479,9 +1481,22 @@ def _ensure_wt_columns(out): ) wt_join["wt_predictor_version"] = wt_join.get("predictor_version") + scope_column = "_topiary_wt_mhc_scope" + support = self.kind_support + + def scope(row): + declared = support.get(row.get(_MODEL_KEY_COLUMN), {}).get(row["kind"], {}) + dependence = declared.get("mhc_dependence") or mhc_dependence(row["kind"], rows=row.to_frame().T) + return prediction_mhc_scope(row.get("allele"), dependence=dependence, allele_set=row.get("allele_set")) + + df = df.copy() + df[scope_column] = df.apply(scope, axis=1) + wt_join[scope_column] = wt_join.apply(scope, axis=1) if not wt_join.empty else pd.Series(dtype=object) + wt_join = wt_join[wt_join[scope_column].notna()] + join_cols = [ _MODEL_KEY_COLUMN, "fragment_id", _WT_OFFSET_COLUMN, - "wt_peptide", "allele", "wt_peptide_length", "kind", + "wt_peptide", scope_column, "wt_peptide_length", "kind", "prediction_method_name", "predictor_version", ] join_cols = [ @@ -1520,7 +1535,7 @@ def _ensure_wt_columns(out): how="left", validate="many_to_one", ) - return _ensure_wt_columns(out) + return _ensure_wt_columns(out.drop(columns=[scope_column])) def _filter_wt_rows_to_baseline_context(self, wt_join, baseline_by_fragment_id): """Keep cached/context rows compatible with the WT baseline flank.""" diff --git a/topiary/ranking/__init__.py b/topiary/ranking/__init__.py index d8a5617..26c1cfb 100644 --- a/topiary/ranking/__init__.py +++ b/topiary/ranking/__init__.py @@ -53,6 +53,7 @@ KIND_MHC_DEPENDENCE, MHC_DEPENDENCE_VALUES, mhc_dependence, + prediction_mhc_scope, resolve_kind, resolve_default_methods, describe_default_versions, @@ -152,6 +153,7 @@ "KIND_MHC_DEPENDENCE", "MHC_DEPENDENCE_VALUES", "mhc_dependence", + "prediction_mhc_scope", "resolve_kind", "resolve_default_methods", "describe_default_versions", diff --git a/topiary/ranking/nodes.py b/topiary/ranking/nodes.py index 2df6506..b26203a 100644 --- a/topiary/ranking/nodes.py +++ b/topiary/ranking/nodes.py @@ -46,6 +46,7 @@ import numpy as np from packaging.version import InvalidVersion, Version import pandas as pd +import mhcgnomes from mhctools import MHC_DEPENDENCE_VALUES, Kind @@ -244,6 +245,46 @@ def format_allele_set(alleles): return _ALLELE_SET_SEPARATOR.join(names) +def prediction_mhc_scope(allele, *, dependence, allele_set=None): + """Return a canonical MHC identity for matching prediction measurements. + + Parameters + ---------- + allele : str or None + The measured allele for per-allele predictions. A haplotype's reported + presenter is deliberately not its measurement identity. + dependence : str + ``single_allele``, ``haplotype`` or ``none``, as resolved by + :func:`mhc_dependence` or declared by the configured model. + allele_set : str or sequence of str, optional + Full genotype for haplotype measurements, as comma-separated names or + a sequence. Alleles are parsed with mhcgnomes and sorted/deduplicated. + + Returns + ------- + tuple or None + Dependence plus canonical allele(s), or ``('none',)`` for allele-free + measurements. Missing required allele/genotype returns None: it is + unmatchable, not an allele-free scope. Exclude missing comparator keys + before joining; pandas otherwise matches null keys to one another. + The tuple identifies only MHC scope; callers must also match occurrence, + kind and model identity. Different haplotype presenters can share a + scope; different genotypes cannot. + """ + if dependence == "none": + return ("none",) + if dependence == "single_allele": + return (dependence, mhcgnomes.parse(str(allele)).to_string()) if is_stated(allele) else None + if dependence != "haplotype": + raise ValueError(f"Unknown MHC dependence: {dependence!r}") + if isinstance(allele_set, (list, tuple, set, frozenset)): + names = allele_set + else: + names = split_allele_set(allele_set) if is_stated(allele_set) else [] + names = sorted({mhcgnomes.parse(str(name)).to_string() for name in names}) + return (dependence, *names) if names else None + + def _peptide_keys(group_keys): """The group keys that identify the peptide rather than the alleles.""" return [k for k in group_keys if k not in _ALLELE_DIMENSION_KEYS] From af5c61f1d260bc64020c425979d535dac78b6ecf Mon Sep 17 00:00:00 2001 From: Alex Rubinsteyn Date: Fri, 2 Oct 2026 00:05:07 -0400 Subject: [PATCH 3/4] Retain compound window identity through occurrence prediction --- docs/api.md | 2 +- docs/combined-sources.md | 4 ++- scripts/check_vaxrank_candidates.py | 41 ++++++++++++++++++++++++++ tests/test_peptide_occurrences.py | 40 ++++++++++++++++++++++++- tests/test_twin_conformance.py | 3 +- topiary/peptide_occurrences.py | 45 ++++++++++++++++++++--------- topiary/predictor.py | 8 +++-- 7 files changed, 122 insertions(+), 21 deletions(-) diff --git a/docs/api.md b/docs/api.md index f9ad333..d0d6a01 100644 --- a/docs/api.md +++ b/docs/api.md @@ -91,7 +91,7 @@ return fresh read-only int8 arrays, so callers never share mutable matrix state. |--------|-------|----------| | `predict_from_named_sequences(dict)` | `{name: sequence}` | Sliding-window scan | | `predict_from_named_peptides(dict)` | `{name: peptide}` | Score as-is | -| `predict_from_peptide_occurrences(table)` | Unique `prediction_id`, peptide, optional flanks/coordinates/annotations | Score exact occurrences with their context; no scanning. | +| `predict_from_peptide_occurrences(table)` | `prediction_id`, peptide, optional flanks/coordinates/annotations; unique by ID, peptide, offset and sample | Score exact occurrences with their context; no scanning. | | `predict_from_sequences(list)` | `[sequence, ...]` | Sliding-window scan | | `predict_from_fragments(fragments)` | `[ProteinFragment]` | Universal path — any origin; fragment-level metadata and `target_intervals` threaded through. | | `predict_from_variants(variants)` | VariantCollection | Variant pipeline (builds `ProteinFragment`s internally and delegates). | diff --git a/docs/combined-sources.md b/docs/combined-sources.md index e2ba2dd..c15680e 100644 --- a/docs/combined-sources.md +++ b/docs/combined-sources.md @@ -538,7 +538,9 @@ assembly. Topiary 5.91.0 adds `predict_peptide_occurrences` and `TopiaryPredictor.predict_from_peptide_occurrences` for a selected peptide -universe. Supply one record per unique `prediction_id`, using the same occurrence +universe. Supply one record per `(prediction_id, peptide, peptide_offset)` within +each sample (when `sample_name` or `candidate_sample` is supplied). An ID may name +several peptide windows in the same source. Use the same occurrence identity Vaxrank and the policy evaluator consume. Different occurrences of one peptide remain separate even when they share a gene, coordinate or HLA candidate. diff --git a/scripts/check_vaxrank_candidates.py b/scripts/check_vaxrank_candidates.py index 4652f9f..12c1922 100644 --- a/scripts/check_vaxrank_candidates.py +++ b/scripts/check_vaxrank_candidates.py @@ -22,6 +22,47 @@ from vaxrank.vaccine_peptide import VaccinePeptide +def test_explicit_windows_keep_repeated_source_ids_through_vaxrank_scoring(): + from topiary import TopiaryPredictor, TopiaryResult, evaluate_selection_policy + from vaxrank.epitope_dsl import prediction_group_columns + + sequence = "AAASIINFEKLTSIINFEKLTTGILGFVFTL" + windows = pd.DataFrame(dict( + prediction_id=["one-protein"] * 3, + peptide=["SIINFEKL", "SIINFEKL", "GILGFVFTL"], + peptide_offset=[3, 12, 22], n_flank=["AAA", "T", "TT"], c_flank=["T", "TT", ""], + source_sequence=[sequence] * 3, + )) + observed = [] + for use_flanks, expected in ((True, [803., 801., 12.]), (False, [800., 800., 10.])): + frame = TopiaryPredictor(models=Model()).predict_from_peptide_occurrences(windows, use_flanks=use_flanks) + assert frame.value.tolist() == expected + assert frame.prediction_id.tolist() == windows.prediction_id.tolist() + rows = [dict( + peptide=row.peptide, source=row.source_sequence, offset=row.peptide_offset, + prediction_id=row.prediction_id, n_flank=row.n_flank, c_flank=row.c_flank, + mutant=Prediction(kind=row.kind, peptide=row.peptide, allele=row.allele, + value=row.value, score=row.score, + predictor_name=row.prediction_method_name, predictor_version=row.predictor_version), + source_class="mutation", overlaps_targetable=True, patient_alleles=[row.allele], + ) for row in frame.itertuples()] + epitopes = candidate_epitopes_from_rows(rows) + policy = SelectionPolicy("explicit-windows", "1 / affinity.value") + cfg = EpitopeConfig(score_expr=policy.score_by, min_epitope_score=0.) + scored = attach_per_allele_scores(epitopes, cfg, topiary_df=frame) + keys = prediction_group_columns(frame) + evaluated = evaluate_selection_policy(TopiaryResult(frame), policy, group_keys=keys) + scores = evaluated.occurrences.set_index(keys).score.to_dict() + assert len(scored) == 3 + for epitope, value in zip(scored, expected): + assert epitope.prediction_id == "one-protein" + score = epitope.per_allele_scores["HLA-A*02:01"] + assert score == pytest.approx(1 / value) + assert score == scores[(*epitope.prediction_group_key, "HLA-A*02:01")] + observed.append([epitope.per_allele_scores for epitope in scored]) + assert observed[0] != observed[1] + + @pytest.mark.parametrize("antigen_kind", ["mutation", "fusion", "splice", "CTA", "ERV", "viral"]) @pytest.mark.parametrize("policy", ["original", "rescored", "rna_overlay"]) def test_candidate_features_reach_vaxrank_scoring_and_vaccine_construction(antigen_kind, policy, tmp_path): diff --git a/tests/test_peptide_occurrences.py b/tests/test_peptide_occurrences.py index 99400a5..dcd42b2 100644 --- a/tests/test_peptide_occurrences.py +++ b/tests/test_peptide_occurrences.py @@ -59,7 +59,7 @@ def test_real_mhctools_named_and_occurrence_doors_agree(peptides): @pytest.mark.parametrize("change,match", [ - ({"prediction_id": "same"}, "unique"), + ({"prediction_id": "same", "peptide": "SIINFEKL", "peptide_offset": 0}, "unique"), ({"prediction_id": None}, "prediction_id"), ({"peptide": ""}, "peptide"), ({"peptide_offset": -1}, "peptide_offset"), @@ -216,6 +216,44 @@ def test_shared_sequences_in_distinct_samples_remain_distinct_inference_contexts pd.testing.assert_frame_equal(output[inputs.columns], inputs) +@pytest.mark.parametrize("predict", PEPTIDE_OCCURRENCE_TWINS) +def test_repeated_source_ids_keep_window_sample_order_and_comparator_identity(predict): + inputs = occurrences().assign(prediction_id="same-source", sample_name="patient-one") + inputs = pd.concat([inputs, inputs.iloc[:1].assign(sample_name="patient-two", n_flank="TT")], ignore_index=True) + inputs["wt_peptide"] = "GILGFVFTL" + inputs["wt_n_flank"] = ["T" * length for length in range(1, 6)] + inputs["wt_c_flank"] = "" + # Both public doors retain the compound source/window/sample identity. + output = predict(inputs, Model(), predict_wt=True) + pd.testing.assert_frame_equal(output[inputs.columns], inputs) + assert output.value.tolist() == [803., 803., 801., 13., 802.] + assert output.wt_value.tolist() == [11., 12., 13., 14., 15.] + filtered = TopiaryPredictor(models=Model(), filter_by="affinity.value < 803").predict_from_peptide_occurrences(inputs) + assert filtered.gene.tolist() == ["G3", "G4", "G1"] + assert filtered.sample_name.tolist() == ["patient-one", "patient-one", "patient-two"] + + +@pytest.mark.parametrize("predict", PEPTIDE_OCCURRENCE_TWINS) +def test_repeated_source_id_with_unknown_coordinates_uses_peptide_identity(predict): + inputs = occurrences().iloc[[0, 3]].drop(columns="peptide_offset").assign(prediction_id="source") + output = predict(inputs, Model()) + pd.testing.assert_frame_equal(output[inputs.columns], inputs.reset_index(drop=True)) + assert output.peptide_offset.isna().all() + assert output.value.tolist() == [803., 13.] + + +def test_partial_cache_reports_each_window_even_when_source_id_repeats(): + cache = CachedPredictor(source().df.iloc[:1].assign(n_flank="", c_flank="", peptide_length=8)) + inputs = occurrences().assign(prediction_id="source") + reports = [] + with pytest.warns(PartialPredictionWarning, match="skipped 1"): + output = predict_peptide_occurrences(inputs, cache, use_flanks=False, on_miss=reports.append) + assert output.peptide_offset.tolist() == [3, 30, 1] + assert reports[0]["source_sequence_name"] == reports[0]["prediction_id"] == "source" + assert reports[0]["peptide"] == "GILGFVFTL" + assert reports[0]["peptide_offset"] == 3 + + @pytest.mark.parametrize("dependence", ["single_allele", "haplotype", "none"]) def test_fragment_and_exact_comparators_match_the_same_mhc_scope(dependence): class ScopeModel: diff --git a/tests/test_twin_conformance.py b/tests/test_twin_conformance.py index 6a5b77e..c54775f 100644 --- a/tests/test_twin_conformance.py +++ b/tests/test_twin_conformance.py @@ -159,7 +159,8 @@ def _sv_support_record(support): # Driven by test_peptide_occurrences.py; source coordinates differ legitimately. PEPTIDE_OCCURRENCE_TWINS = ( predict_peptide_occurrences, - lambda occurrences, model, **kwargs: TopiaryPredictor(models=model).predict_from_peptide_occurrences(occurrences, **kwargs), + lambda occurrences, model, predict_wt=False, **kwargs: TopiaryPredictor( + models=model, predict_wt=predict_wt).predict_from_peptide_occurrences(occurrences, **kwargs), ) NAMED_OCCURRENCE_TWINS = ( lambda names, model: TopiaryPredictor(models=model).predict_from_named_peptides(names), diff --git a/topiary/peptide_occurrences.py b/topiary/peptide_occurrences.py index 74928e8..bac5d5c 100644 --- a/topiary/peptide_occurrences.py +++ b/topiary/peptide_occurrences.py @@ -16,7 +16,10 @@ def predict_peptide_occurrences(occurrences, model, *, use_flanks=True, predict_ Parameters ---------- occurrences : pandas.DataFrame or iterable of mappings - One row per unique, nonempty string ``prediction_id`` and ``peptide``. + Nonempty string ``prediction_id`` and ``peptide`` identify each row + together with ``peptide_offset`` and any supplied ``sample_name`` / + ``candidate_sample``. This compound identity must be unique; an ID may + name several peptide windows in one source. Optional ``peptide_offset`` is a nonnegative source coordinate; absent coordinates remain unknown. ``n_flank`` and ``c_flank`` are strings or missing: an empty string states a known terminus. Other non-prediction @@ -42,7 +45,8 @@ def predict_peptide_occurrences(occurrences, model, *, use_flanks=True, predict_ on_miss : callable, optional Explicitly report and skip cache coverage failures, using :func:`predict_with_cache_miss_report`. The report's source name is the - occurrence ID. Other errors still raise; partial results emit a warning. + occurrence ID, with peptide, offset and supplied sample keys identifying + the window. Other errors still raise; partial results emit a warning. Returns ------- @@ -67,8 +71,6 @@ def predict_peptide_occurrences(occurrences, model, *, use_flanks=True, predict_ for column in ("prediction_id", "peptide"): if column not in frame or not frame[column].map(lambda value: isinstance(value, str) and bool(value.strip())).all(): raise ValueError(f"Occurrences require nonempty string {column}") - if frame.prediction_id.duplicated().any(): - raise ValueError("prediction_id must be unique; retain distinct source observations under distinct IDs") reserved = (PREDICTION_COLUMNS | {"allele", "value_unit", "measurement_context", "predictor_name", "offset", "prediction_mhc_dependence", "prediction_flanks_supplied", "candidate_id", "candidate_allele", "candidate_mhc_class"}) & set(frame) @@ -97,6 +99,10 @@ def predict_peptide_occurrences(occurrences, model, *, use_flanks=True, predict_ and isinstance(value, (int, float, np.integer, np.floating)) and np.isfinite(value) and value >= 0 and value == int(value)).all(): raise ValueError("peptide_offset must be a nonnegative integer or missing") + identity_columns = ["prediction_id", "peptide", "peptide_offset"] + [ + column for column in ("sample_name", "candidate_sample") if column in frame] + if frame.duplicated(identity_columns).any(): + raise ValueError("Occurrence identity must be unique by prediction_id, peptide, peptide_offset and sample") if "peptide_length" in frame and not frame.peptide_length.eq(frame.peptide.str.len()).all(): raise ValueError("peptide_length disagrees with peptide") flank_model = bool(getattr(model, "uses_flanking_sequences", False)) @@ -118,11 +124,16 @@ def predict_peptide_occurrences(occurrences, model, *, use_flanks=True, predict_ if any(genotype and genotype != format_allele_set(alleles) for genotype in genotypes): raise ValueError("Haplotype occurrence allele_set must match the configured model") if on_miss is not None: + def report_miss(report): + row = frame.iloc[int(report["source_sequence_name"])] + identity = {column: None if pd.isna(row[column]) else row[column] for column in identity_columns} + on_miss(dict(report, **identity, source_sequence_name=row.prediction_id)) + reported = predict_with_cache_miss_report( lambda inputs: predict_peptide_occurrences( - frame.loc[frame.prediction_id.isin(inputs)], model, + frame.iloc[[int(index) for index in inputs]], model, use_flanks=use_flanks, predict_wt=predict_wt), - dict(zip(frame.prediction_id, frame.peptide)), on_miss=on_miss, + {str(index): peptide for index, peptide in frame.peptide.items()}, on_miss=report_miss, context=dict(stage="occurrence", prediction_method_name=str( getattr(model, "prediction_method_name", getattr(model, "predictor_name", type(model).__name__))), predictor_version=str(getattr(model, "predictor_version", "")), @@ -193,18 +204,20 @@ def predict_peptide_occurrences(occurrences, model, *, use_flanks=True, predict_ context = inputs.drop(columns=["allele_set"], errors="ignore") expanded.append(context.merge(payload, on="peptide", how="left", validate="many_to_many")) output = _concat_frames(expanded) - order = {name: i for i, name in enumerate(frame.prediction_id)} - output = output.iloc[np.argsort(output.prediction_id.map(order), kind="stable")].reset_index(drop=True) + positions = pd.MultiIndex.from_frame(frame[identity_columns]).get_indexer( + pd.MultiIndex.from_frame(output[identity_columns])) + order = np.argsort(positions, kind="stable") + output = output.iloc[order].reset_index(drop=True) if "source_sequence_name" not in output: output["source_sequence_name"] = output.prediction_id if "sample_name" not in output: output["sample_name"] = "" if predict_wt: - output = _attach_comparators(output, frame, model, use_flanks) + output = _attach_comparators(output, frame, positions[order], model, use_flanks) return output -def _attach_comparators(output, occurrences, model, use_flanks): +def _attach_comparators(output, occurrences, positions, model, use_flanks): """Plumbing for the explicit comparator pass of predict_peptide_occurrences.""" present = occurrences.get("wt_peptide", pd.Series(None, index=occurrences.index, dtype=object)).map( lambda value: isinstance(value, str) and bool(value)) @@ -214,6 +227,9 @@ def _attach_comparators(output, occurrences, model, use_flanks): output["wt_" + field] = np.nan return output comparators = occurrences.loc[present].copy() + # Comparator sequence/offset need not match the primary occurrence. Use + # its original row position for this private pass, retaining public IDs. + comparators["prediction_id"] = comparators.index.astype(str) comparators["peptide"] = comparators.wt_peptide comparators = comparators.drop(columns=["peptide_length"], errors="ignore") for column in ("n_flank", "c_flank", "peptide_offset"): @@ -222,11 +238,12 @@ def _attach_comparators(output, occurrences, model, use_flanks): scores = predict_peptide_occurrences(comparators, model, use_flanks=use_flanks) # A haplotype's deconvolved presenter may change for the comparator. Its # score still describes the same configured genotype, not that one allele. - def key(row): - return (row.prediction_id, row.kind, prediction_mhc_scope( + def key(identity, row): + return (identity, row.kind, prediction_mhc_scope( row.allele, dependence=row.prediction_mhc_dependence, allele_set=row.allele_set)) - by_scope = {key(row): row for row in scores.itertuples(index=False)} - matches = [by_scope.get(key(row)) for row in output.itertuples(index=False)] + by_scope = {key(row.prediction_id, row): row for row in scores.itertuples(index=False)} + matches = [by_scope.get(key(str(position), row)) + for position, row in zip(positions, output.itertuples(index=False))] for field in fields: output["wt_" + field] = [getattr(row, field, np.nan) for row in matches] return output diff --git a/topiary/predictor.py b/topiary/predictor.py index edf7ce2..20b1cee 100644 --- a/topiary/predictor.py +++ b/topiary/predictor.py @@ -1013,8 +1013,9 @@ def predict_from_peptide_occurrences(self, occurrences, *, use_flanks=True): Parameters ---------- occurrences : pandas.DataFrame or iterable of mappings - Explicit unique ``prediction_id`` and peptide, with optional source - coordinates, flanks and non-prediction annotations. See + Explicit ``prediction_id``, peptide and optional source coordinates, + flanks and non-prediction annotations. Identity is unique by ID, + peptide, offset and supplied sample keys. See :func:`predict_peptide_occurrences` for validation and scope rules. use_flanks : bool Use supplied flanks. False explicitly requests peptide-only scores, @@ -1040,7 +1041,8 @@ def predict_from_peptide_occurrences(self, occurrences, *, use_flanks=True): occurrences, model, use_flanks=use_flanks, predict_wt=self.predict_wt, on_miss=handler), model_key)) df = _concat_frames(frames) if frames else pd.DataFrame() - keys = ["prediction_id", "peptide", "peptide_offset", "allele"] + keys = ["prediction_id", "peptide", "peptide_offset", "allele"] + [ + column for column in ("sample_name", "candidate_sample") if column in occurrences] return self._attach_result_attrs(self._strip_internal_columns( self._apply_filter(df, group_keys=keys)).reset_index(drop=True)) From 354e33aa97d0b3ba6570c8f943464e3cc96606cc Mon Sep 17 00:00:00 2001 From: Alex Rubinsteyn Date: Fri, 2 Oct 2026 00:22:00 -0400 Subject: [PATCH 4/4] Build mixed-kind fixture consistently in both pandas string modes --- tests/test_candidate_tables.py | 8 +++++--- tests/test_peptide_occurrences.py | 2 +- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/test_candidate_tables.py b/tests/test_candidate_tables.py index 1310cb7..d3a7876 100644 --- a/tests/test_candidate_tables.py +++ b/tests/test_candidate_tables.py @@ -394,15 +394,17 @@ def kind_support(self): def predict_dataframe(self, peptides, **kwargs): base = super().predict_dataframe(peptides, **kwargs) - frames = [] + records = [] for kind in self.emitted: frame = base.copy() frame["kind"] = kind if kind == "antigen_processing": frame["allele"] = None frame["value_unit"] = None - frames.append(frame) - return pd.concat(frames, ignore_index=True) + records.extend(frame.to_dict("records")) + # Infer one dtype across both allele-bearing and allele-free rows, + # including pandas' inferred-string mode (#454). + return pd.DataFrame(records) @pytest.mark.parametrize("alleles", [["HLA-B*07:02"] * 2, ["HLA-A*02:01", "HLA-B*07:02"]]) diff --git a/tests/test_peptide_occurrences.py b/tests/test_peptide_occurrences.py index dcd42b2..7218fb7 100644 --- a/tests/test_peptide_occurrences.py +++ b/tests/test_peptide_occurrences.py @@ -95,7 +95,7 @@ def test_unknown_and_empty_flanks_are_distinct_and_positions_not_invented(tmp_pa predict_peptide_occurrences(inputs, Model()) -def test_model_scope_and_coverage_are_explicit(): +def test_model_scope_and_coverage_are_explicit(pandas_string_inference): inputs = occurrences() kinds = ("pMHC_affinity", "antigen_processing") output = predict_peptide_occurrences(inputs, MultiKindModel(declared=kinds, emitted=kinds))