Add predictor MHC context metadata

Author: iskandr

README.md

@@ -152,10 +152,12 @@ df = ms.predict_dataframe(["SIINFEKL"])
 df = ms.predict_proteins_dataframe({"TP53": "MEEPQ..."})
 ```
 
-### Measurement kinds
+### Measurement kinds and MHC context
 
 Each `Prediction` has a `kind` string describing what it measures:
 
+The canonical prediction kind strings are defined in `mhctools.pred.Kind`.
+
 | Kind | Meaning |
 |---|---|
 | `pMHC_affinity` | Peptide-MHC binding affinity |
@@ -167,6 +169,43 @@ Each `Prediction` has a `kind` string describing what it measures:
 | `tap_transport` | TAP transport score (reserved, not yet used) |
 | `erap_trimming` | ERAP trimming score (reserved, not yet used) |
 
+Predictors also expose `kind_support()` so downstream code can tell what MHC
+context is meaningful for each emitted kind:
+
+```python
+support = predictor.kind_support()
+support["pMHC_affinity"]
+# {"mhc_dependence": "single_allele", "mhc_class": "I"}
+```
+
+`mhc_dependence` is one of:
+
+| Value | Meaning |
+|---|---|
+| `none` | The prediction is MHC-independent; `Prediction.allele` is empty. |
+| `single_allele` | The prediction is for one peptide/MHC allele pair; `Prediction.allele` is part of the key. |
+| `haplotype` | The prediction uses the requested MHC repertoire jointly; `Prediction.allele` may carry best-allele attribution but is not the prediction key. |
+
+`mhc_class` is one of `none`, `I`, `II`, or `both`.
+
+The allowed metadata values are defined in `mhctools.pred` as
+`MHC_DEPENDENCE_VALUES` and `MHC_CLASS_VALUES`.
+
+Examples:
+
+| Predictor | Kind | `mhc_dependence` | `mhc_class` |
+|---|---|---|---|
+| `NetMHCpan41` | `pMHC_affinity` | `single_allele` | `I` |
+| `NetMHCpan41` | `pMHC_presentation` | `single_allele` | `I` |
+| `MHCflurry` | `pMHC_affinity` | `single_allele` | `I` |
+| `MHCflurry` | `pMHC_presentation` | `haplotype` | `I` |
+| `Pepsickle` | `proteasome_cleavage` | `none` | `none` |
+
+For MHCflurry presentation, mhctools calls the upstream presentation model with
+the full requested allele set and emits one `pMHC_presentation` record per
+peptide. The `allele` field carries MHCflurry's `best_allele` attribution when
+available.
+
 ### The Prediction object
 
 Every prediction is a frozen, self-contained `Prediction` dataclass:

mhctools/__init__.py

@@ -1,6 +1,15 @@
 from .binding_prediction import BindingPrediction
 from .binding_prediction_collection import BindingPredictionCollection
-from .pred import Prediction, Pred, PeptideResult, PeptidePreds, Kind, preds_from_rows
+from .pred import (
+    Kind,
+    MHC_CLASS_VALUES,
+    MHC_DEPENDENCE_VALUES,
+    PeptidePreds,
+    PeptideResult,
+    Pred,
+    Prediction,
+    preds_from_rows,
+)
 from .sample import MultiSample
 from .iedb import (
     IedbNetMHCcons,
@@ -63,14 +72,16 @@ def __getattr__(name):
     raise AttributeError(
         "module %r has no attribute %r" % (__name__, name))
 
-__version__ = "3.13.6"
+__version__ = "3.13.7"
 
 __all__ = [
     "Prediction",
     "Pred",  # backward compat alias
     "PeptideResult",
     "PeptidePreds",  # backward compat alias
     "Kind",
+    "MHC_CLASS_VALUES",
+    "MHC_DEPENDENCE_VALUES",
     "preds_from_rows",
     "MultiSample",
     "BindingPrediction",

mhctools/base_predictor.py

@@ -19,7 +19,11 @@
 
 from .unsupported_allele import UnsupportedAllele
 from .binding_prediction_collection import BindingPredictionCollection
-from .pred import Prediction, Kind, PeptideResult
+from .pred import (
+    Kind,
+    PeptideResult,
+    Prediction,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -101,6 +105,7 @@ class BasePredictor(object):
     flank_length = 15
     n_flank_length = None
     c_flank_length = None
+    mhc_class = "I"
 
     def __init__(
             self,
@@ -331,6 +336,20 @@ def _default_pred_kind(self):
         """Override in subclasses to set the Kind for compat conversion."""
         return Kind.pMHC_affinity
 
+    def kind_support(self):
+        """Predictor-specific MHC context for supported prediction kinds."""
+        return {
+            self._default_pred_kind(): {
+                "mhc_dependence": "single_allele",
+                "mhc_class": self.mhc_class,
+            }
+        }
+
+    @property
+    def supported_kinds(self):
+        """Prediction kind strings this predictor can emit."""
+        return tuple(self.kind_support())
+
     # --- deprecated API (still works) ---
 
     def predict_peptides(self, peptides):

mhctools/bigmhc.py

@@ -31,7 +31,12 @@
 import pandas as pd
 import torch
 
-from .pred import Kind, Prediction, PeptideResult, COLUMNS
+from .pred import (
+    COLUMNS,
+    Kind,
+    PeptideResult,
+    Prediction,
+)
 
 
 def _find_bigmhc_dir(bigmhc_path=None):
@@ -133,6 +138,18 @@ def _pred_kind(self):
             return Kind.immunogenicity
         return Kind.pMHC_presentation
 
+    def kind_support(self):
+        return {
+            self._pred_kind(): {
+                "mhc_dependence": "single_allele",
+                "mhc_class": "I",
+            }
+        }
+
+    @property
+    def supported_kinds(self):
+        return tuple(self.kind_support())
+
     def _predictor_name(self):
         return "bigmhc_%s" % self.mode
 

mhctools/iedb.py

@@ -325,6 +325,8 @@ def __init__(
 IEDB_MHC_CLASS_II_URL = "http://tools-cluster-interface.iedb.org/tools_api/mhcii/"
 
 class IedbNetMHCIIpan(IedbBasePredictor):
+    mhc_class = "II"
+
     def __init__(
             self,
             alleles,

mhctools/mhcflurry.py

@@ -18,7 +18,7 @@
 from .base_predictor import _check_flank_inputs
 from .binding_prediction import BindingPrediction
 from .binding_prediction_collection import BindingPredictionCollection
-from .pred import Prediction, Kind
+from .pred import Kind, Prediction
 from .unsupported_allele import UnsupportedAllele
 
 logger = logging.getLogger(__name__)
@@ -99,9 +99,10 @@ class MHCflurry(BasePredictor):
     """
     MHCflurry predictor using the modern Class1PresentationPredictor API.
 
-    Produces both ``pMHC_affinity`` and ``pMHC_presentation`` predictions
-    per peptide-allele pair. The legacy ``predict_peptides`` method returns
-    BindingPrediction objects based on affinity values for backward compat.
+    Produces per-allele ``pMHC_affinity`` predictions and one haplotype-level
+    ``pMHC_presentation`` prediction per peptide. The legacy
+    ``predict_peptides`` method returns BindingPrediction objects based on
+    affinity values for backward compat.
 
     See https://github.com/openvax/mhcflurry
     """
@@ -214,11 +215,12 @@ def predict(self, peptides, n_flanks=None, c_flanks=None):
         """
         Predict for a list of peptide sequences.
 
-        Returns a list of PeptideResult, each containing both
-        pMHC_affinity and pMHC_presentation Prediction objects per allele.
+        Returns a list of PeptideResult, each containing one pMHC_affinity
+        Prediction per allele and one haplotype-level pMHC_presentation
+        Prediction with best_allele attribution when MHCflurry reports it.
 
-        Uses batch prediction across all alleles in a single call for
-        both affinity and presentation scores.
+        Uses batch prediction across alleles for affinity and a genotype-level
+        presentation call for presentation scores.
         """
         from .pred import PeptideResult
 
@@ -243,34 +245,38 @@ def predict(self, peptides, n_flanks=None, c_flanks=None):
             include_percentile_ranks=self.include_affinity_percentile_ranks,
         )
 
-        # Per-allele presentation calls (presentation predictor does
-        # deconvolution across alleles, so we call per-allele to get
-        # per-allele presentation scores). Key by mhcflurry's output
-        # allele string so lookups with aff_df.allele always match.
-        pres_by_pep_allele = {}
-        for input_allele in allele_list:
-            kwargs = {
-                "peptides": peptide_list,
-                "alleles": [input_allele],
-                "include_affinity_percentile": False,
-                "verbose": 0,
-            }
-            if n_flank_list is not None:
-                kwargs["n_flanks"] = n_flank_list
-            if c_flank_list is not None:
-                kwargs["c_flanks"] = c_flank_list
-            df = self.predictor.predict(**kwargs)
-            if len(df) != len(peptide_list):
+        # The presentation predictor deconvolves across the full genotype.
+        # Preserve that haplotype-level contract here: one presentation
+        # Prediction per peptide, with best_allele carried as attribution
+        # when MHCflurry reports it.
+        kwargs = {
+            "peptides": peptide_list,
+            "alleles": allele_list,
+            "include_affinity_percentile": False,
+            "verbose": 0,
+        }
+        if n_flank_list is not None:
+            kwargs["n_flanks"] = n_flank_list
+        if c_flank_list is not None:
+            kwargs["c_flanks"] = c_flank_list
+        pres_df = self.predictor.predict(**kwargs)
+        if len(pres_df) != len(peptide_list):
+            raise ValueError(
+                "MHCflurry returned %d presentation row(s) for %d "
+                "peptide input(s)" % (len(pres_df), len(peptide_list)))
+
+        pres_by_peptide_index = {}
+        for row_position, row in enumerate(pres_df.itertuples(index=False)):
+            row_index = int(getattr(row, "peptide_num", row_position))
+            if row_index in pres_by_peptide_index:
                 raise ValueError(
-                    "MHCflurry returned %d presentation row(s) for %d "
-                    "peptide input(s) and allele '%s'" % (
-                        len(df), len(peptide_list), input_allele))
-            for row_index, row in enumerate(df.itertuples(index=False)):
-                output_allele = getattr(row, 'allele', input_allele)
-                pres_by_pep_allele[(row_index, output_allele)] = (
-                    row.presentation_score,
-                    row.presentation_percentile,
-                )
+                    "MHCflurry returned duplicate presentation row for "
+                    "peptide index %d" % row_index)
+            pres_by_peptide_index[row_index] = (
+                row.presentation_score,
+                row.presentation_percentile,
+                getattr(row, "best_allele", getattr(row, "allele", "")),
+            )
 
         groups = [list() for _ in peptide_list]
         for row_index, row in zip(batch_indices, aff_df.itertuples(index=False)):
@@ -301,19 +307,21 @@ def predict(self, peptides, n_flanks=None, c_flanks=None):
                 predictor_name="mhcflurry",
             ))
 
-            key = (row_index, allele)
-            if key not in pres_by_pep_allele:
+        for row_index, pep in enumerate(peptide_list):
+            if row_index not in pres_by_peptide_index:
                 raise ValueError(
                     "MHCflurry: missing presentation score for "
-                    "peptide='%s' allele='%s' (this indicates an allele or "
-                    "peptide string mismatch between the affinity and "
-                    "presentation predictor outputs)" % (pep, allele))
-            pres_score, pres_pct = pres_by_pep_allele[key]
+                    "peptide index %d, peptide='%s'" % (row_index, pep))
+            pres_score, pres_pct, best_allele = pres_by_peptide_index[row_index]
+            n_flank = (
+                n_flank_list[row_index] if n_flank_list is not None else "")
+            c_flank = (
+                c_flank_list[row_index] if c_flank_list is not None else "")
             groups[row_index].append(Prediction(
                 kind=Kind.pMHC_presentation,
                 score=pres_score,
                 peptide=pep,
-                allele=allele,
+                allele=best_allele or "",
                 n_flank=n_flank,
                 c_flank=c_flank,
                 percentile_rank=pres_pct,
@@ -331,6 +339,18 @@ def predict_with_flanks(self, peptides, n_flanks, c_flanks):
     def _default_pred_kind(self):
         return Kind.pMHC_affinity
 
+    def kind_support(self):
+        return {
+            Kind.pMHC_affinity: {
+                "mhc_dependence": "single_allele",
+                "mhc_class": "I",
+            },
+            Kind.pMHC_presentation: {
+                "mhc_dependence": "haplotype",
+                "mhc_class": "I",
+            },
+        }
+
 
 class MHCflurry_Affinity(BasePredictor):
     """

mhctools/netmhc_pan4.py

@@ -10,9 +10,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from functools import partial
+
 from .base_commandline_predictor import BaseCommandlinePredictor
 from .parsing import parse_netmhcpan4_stdout, parse_netmhcpan_to_preds
-from functools import partial
+from .pred import Kind
+
 
 class NetMHCpan4(BaseCommandlinePredictor):
     def __init__(
@@ -37,6 +40,7 @@ def __init__(
             flags = []
         else:
             raise ValueError("Unsupported mode", mode)
+        self.mode = mode
 
         BaseCommandlinePredictor.__init__(
             self,
@@ -52,6 +56,18 @@ def __init__(
             extra_flags=flags + extra_flags,
             process_limit=process_limit)
 
+    def kind_support(self):
+        kind = (
+            Kind.pMHC_affinity
+            if self.mode == "binding_affinity" else Kind.pMHC_presentation)
+        return {
+            kind: {
+                "mhc_dependence": "single_allele",
+                "mhc_class": "I",
+            }
+        }
+
+
 class NetMHCpan4_EL(NetMHCpan4):
     """
     Wrapper for NetMHCpan4 when the preferred mode is elution score