From 8314526e3757c986ceacf683dfa7fd5e287efd4c Mon Sep 17 00:00:00 2001 From: Alex Rubinsteyn Date: Fri, 10 Jul 2026 19:33:16 -0400 Subject: [PATCH] Add BigMHC paired prediction API --- mhctools/__init__.py | 2 +- mhctools/bigmhc.py | 97 +++++++++++++++++++++++++++++-------- tests/test_bigmhc_pairs.py | 98 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 177 insertions(+), 20 deletions(-) create mode 100644 tests/test_bigmhc_pairs.py diff --git a/mhctools/__init__.py b/mhctools/__init__.py index 0b54d37..f90ef2c 100644 --- a/mhctools/__init__.py +++ b/mhctools/__init__.py @@ -87,7 +87,7 @@ def __getattr__(name): raise AttributeError( "module %r has no attribute %r" % (__name__, name)) -__version__ = "3.31.1" +__version__ = "3.31.2" __all__ = [ "Prediction", diff --git a/mhctools/bigmhc.py b/mhctools/bigmhc.py index 7d2cf1f..42bc5fc 100644 --- a/mhctools/bigmhc.py +++ b/mhctools/bigmhc.py @@ -32,6 +32,7 @@ import torch from .pred import ( + COLUMNS, Kind, PeptideResult, Prediction, @@ -237,6 +238,77 @@ def _predict_raw(self, peptides, alleles): # Public API # ------------------------------------------------------------------ + def predict_pairs(self, peptides, alleles=None): + """Predict scores for explicit peptide/allele pairs. + + Parameters + ---------- + peptides : list of str or iterable of (str, str) + Peptides to score. If *alleles* is omitted, this must be an + iterable of ``(peptide, allele)`` pairs. + alleles : list of str, optional + Alleles parallel to *peptides*. + + Returns + ------- + list of PeptideResult + One entry per input pair, in order; each contains one Prediction. + """ + if alleles is None: + pairs = list(peptides) + peptide_list = [] + allele_list = [] + for pair in pairs: + try: + peptide, allele = pair + except (TypeError, ValueError): + raise ValueError( + "Expected (peptide, allele) pairs, got %r" % (pair,)) + peptide_list.append(peptide) + allele_list.append(allele) + peptide_list = self._normalize_peptides(peptide_list) + else: + peptide_list = self._normalize_peptides(peptides) + if isinstance(alleles, str): + allele_list = [alleles] + else: + allele_list = list(alleles) + if len(peptide_list) != len(allele_list): + raise ValueError( + "peptides length %d != alleles length %d" + % (len(peptide_list), len(allele_list))) + + allele_list = [str(allele).strip() for allele in allele_list] + if not peptide_list: + return [] + + scores = self._predict_raw(peptide_list, allele_list) + if len(scores) != len(peptide_list): + raise ValueError( + "BigMHC returned %d scores for %d peptide/allele pairs" + % (len(scores), len(peptide_list))) + + kind = self._default_pred_kind() + name = self._predictor_name() + results = [] + for peptide, allele, score in zip(peptide_list, allele_list, scores): + results.append(PeptideResult(preds=(Prediction( + kind=kind, + score=float(score), + peptide=peptide, + allele=allele, + predictor_name=name, + ),))) + return results + + def predict_pairs_dataframe(self, peptides, alleles=None, sample_name=""): + """``predict_pairs()`` flattened to a DataFrame.""" + dfs = [pp.to_dataframe(sample_name) + for pp in self.predict_pairs(peptides, alleles)] + if not dfs: + return pd.DataFrame(columns=COLUMNS) + return pd.concat(dfs, ignore_index=True) + def predict(self, peptides): """Predict scores for peptides against all alleles. @@ -251,31 +323,18 @@ def predict(self, peptides): """ peptides = self._normalize_peptides(peptides) - # Build all (peptide, allele) combinations - all_peptides = [] - all_alleles = [] + pairs = [] for pep in peptides: for allele in self.alleles: - all_peptides.append(pep) - all_alleles.append(allele) - - scores = self._predict_raw(all_peptides, all_alleles) - - kind = self._default_pred_kind() - name = self._predictor_name() + pairs.append((pep, allele)) + flat_results = self.predict_pairs(pairs) idx = 0 results = [] - for pep in peptides: + for _ in peptides: preds = [] - for allele in self.alleles: - preds.append(Prediction( - kind=kind, - score=float(scores[idx]), - peptide=pep, - allele=allele, - predictor_name=name, - )) + for _ in self.alleles: + preds.extend(flat_results[idx].preds) idx += 1 results.append(PeptideResult(preds=tuple(preds))) return results diff --git a/tests/test_bigmhc_pairs.py b/tests/test_bigmhc_pairs.py new file mode 100644 index 0000000..19ef342 --- /dev/null +++ b/tests/test_bigmhc_pairs.py @@ -0,0 +1,98 @@ +import pytest + +from mhctools.bigmhc import BigMHC +from mhctools.pred import COLUMNS, Kind, PeptideResult + + +class _FakeBigMHC(BigMHC): + def __init__(self, alleles=("HLA-A*02:01",), mode="im"): + self.alleles = list(alleles) + self.mode = mode + self.calls = [] + + def _predict_raw(self, peptides, alleles): + self.calls.append((list(peptides), list(alleles))) + return [i / 10.0 for i in range(len(peptides))] + + +def test_predict_pairs_scores_parallel_peptide_allele_rows(): + predictor = _FakeBigMHC(mode="im") + + results = predictor.predict_pairs( + ["siinfekl", "GILGFVFTL"], + ["HLA-B*07:02", "HLA-A*02:01"]) + + assert predictor.calls == [( + ["SIINFEKL", "GILGFVFTL"], + ["HLA-B*07:02", "HLA-A*02:01"], + )] + assert len(results) == 2 + assert all(isinstance(result, PeptideResult) for result in results) + assert [result.preds[0].peptide for result in results] == [ + "SIINFEKL", "GILGFVFTL"] + assert [result.preds[0].allele for result in results] == [ + "HLA-B*07:02", "HLA-A*02:01"] + assert [result.preds[0].score for result in results] == [0.0, 0.1] + assert results[0].preds[0].kind == Kind.immunogenicity + assert results[0].preds[0].predictor_name == "bigmhc_im" + + +def test_predict_pairs_accepts_pair_iterable(): + predictor = _FakeBigMHC(mode="el") + + results = predictor.predict_pairs([ + ("SIINFEKL", "HLA-B*07:02"), + ("GILGFVFTL", "HLA-A*02:01"), + ]) + + assert predictor.calls == [( + ["SIINFEKL", "GILGFVFTL"], + ["HLA-B*07:02", "HLA-A*02:01"], + )] + assert results[0].preds[0].kind == Kind.pMHC_presentation + assert results[0].preds[0].predictor_name == "bigmhc_el" + + +def test_predict_pairs_rejects_length_mismatch(): + predictor = _FakeBigMHC() + + with pytest.raises(ValueError) as e: + predictor.predict_pairs(["SIINFEKL"], ["HLA-A*02:01", "HLA-B*07:02"]) + + assert "peptides length 1 != alleles length 2" in str(e.value) + assert predictor.calls == [] + + +def test_predict_uses_pairs_for_existing_cross_product_api(): + predictor = _FakeBigMHC( + alleles=["HLA-A*02:01", "HLA-B*07:02"], + mode="el") + + results = predictor.predict(["siinfekl", "GILGFVFTL"]) + + assert predictor.calls == [( + ["SIINFEKL", "SIINFEKL", "GILGFVFTL", "GILGFVFTL"], + ["HLA-A*02:01", "HLA-B*07:02", "HLA-A*02:01", "HLA-B*07:02"], + )] + assert len(results) == 2 + assert [pred.allele for pred in results[0].preds] == [ + "HLA-A*02:01", "HLA-B*07:02"] + assert [pred.score for pred in results[0].preds] == [0.0, 0.1] + assert [pred.allele for pred in results[1].preds] == [ + "HLA-A*02:01", "HLA-B*07:02"] + assert [pred.score for pred in results[1].preds] == [0.2, 0.3] + + +def test_predict_pairs_dataframe_is_one_row_per_pair(): + predictor = _FakeBigMHC(mode="im") + + df = predictor.predict_pairs_dataframe( + ["SIINFEKL", "GILGFVFTL"], + ["HLA-A*02:01", "HLA-B*07:02"], + sample_name="sample-1") + + assert list(df.columns) == list(COLUMNS) + assert len(df) == 2 + assert df["sample_name"].tolist() == ["sample-1", "sample-1"] + assert df["peptide"].tolist() == ["SIINFEKL", "GILGFVFTL"] + assert df["allele"].tolist() == ["HLA-A*02:01", "HLA-B*07:02"]