Add BigMHC paired prediction API

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",

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

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"]