[codex] Check MHCflurry percentile calibration

mhctools/__init__.py

@@ -63,7 +63,7 @@ def __getattr__(name):
     raise AttributeError(
         "module %r has no attribute %r" % (__name__, name))
 
-__version__ = "3.13.3"
+__version__ = "3.13.4"
 
 __all__ = [
     "Prediction",

mhctools/mhcflurry.py

@@ -14,8 +14,6 @@
 import math
 import os
 
-from numpy import nan
-
 from .base_predictor import BasePredictor
 from .binding_prediction import BindingPrediction
 from .binding_prediction_collection import BindingPredictionCollection
@@ -26,6 +24,7 @@
 
 # Module-level cache for loaded models. Keyed by (kind, normalized_path).
 _model_cache = {}
+_PERCENT_RANK_SUPPORT_UNKNOWN = object()
 
 
 def _normalize_models_path(models_path):
@@ -41,6 +40,60 @@ def _normalize_models_path(models_path):
     return os.path.realpath(os.path.expanduser(models_path))
 
 
+def _affinity_percent_rank_calibrated_allele(affinity_predictor, allele):
+    """Return the allele whose affinity percentile calibration can be used."""
+    helper = getattr(
+        affinity_predictor, "percent_rank_calibrated_allele", None)
+    if callable(helper):
+        return helper(allele)
+
+    transforms = getattr(
+        affinity_predictor, "allele_to_percent_rank_transform", None)
+    if transforms is None:
+        return _PERCENT_RANK_SUPPORT_UNKNOWN
+
+    canonicalize = getattr(
+        affinity_predictor, "canonicalize_allele_name", None)
+    normalized_allele = (
+        canonicalize(allele) if callable(canonicalize) else allele)
+    if normalized_allele in transforms:
+        return normalized_allele
+
+    allele_to_sequence = getattr(affinity_predictor, "allele_to_sequence", None)
+    if (
+            not allele_to_sequence
+            or normalized_allele not in allele_to_sequence):
+        return None
+
+    sequence = allele_to_sequence[normalized_allele]
+    for other_allele in sorted(allele_to_sequence):
+        if (
+                allele_to_sequence[other_allele] == sequence
+                and other_allele in transforms):
+            return other_allele
+    return None
+
+
+def _check_affinity_percent_rank_support(affinity_predictor, alleles):
+    """Raise if requested alleles cannot get affinity percentile ranks."""
+    missing_alleles = []
+    for allele in alleles:
+        calibrated = _affinity_percent_rank_calibrated_allele(
+            affinity_predictor, allele)
+        if calibrated is _PERCENT_RANK_SUPPORT_UNKNOWN:
+            return
+        if calibrated is None:
+            missing_alleles.append(allele)
+    if missing_alleles:
+        raise ValueError(
+            "MHCflurry affinity percentile ranks are unavailable for "
+            "allele(s): %s. Raw affinity prediction may still be supported. "
+            "Pass include_affinity_percentile_ranks=False to omit affinity "
+            "percentile ranks, or calibrate MHCflurry percentile ranks for "
+            "these alleles."
+            % ", ".join(sorted(missing_alleles)))
+
+
 class MHCflurry(BasePredictor):
     """
     MHCflurry predictor using the modern Class1PresentationPredictor API.
@@ -57,7 +110,8 @@ def __init__(
             alleles,
             default_peptide_lengths=[9],
             predictor=None,
-            models_path=None):
+            models_path=None,
+            include_affinity_percentile_ranks=True):
         """
         Parameters
         -----------
@@ -70,6 +124,12 @@ def __init__(
 
         models_path : string
             Models dir to use if predictor argument is None
+
+        include_affinity_percentile_ranks : bool
+            Whether to request affinity percentile ranks. Enabled by default.
+            If enabled, requested alleles must have MHCflurry affinity
+            percentile-rank calibration, either directly or through an allele
+            with the same pseudosequence.
         """
         from mhcflurry import Class1PresentationPredictor
         BasePredictor.__init__(
@@ -93,9 +153,15 @@ def __init__(
                         Class1PresentationPredictor.load()
             self.predictor = _model_cache[cache_key]
 
+        self.include_affinity_percentile_ranks = \
+            include_affinity_percentile_ranks
+
         for allele in self.alleles:
             if allele not in self.predictor.supported_alleles:
                 raise UnsupportedAllele(allele)
+        if self.include_affinity_percentile_ranks:
+            _check_affinity_percent_rank_support(
+                self.predictor.affinity_predictor, self.alleles)
 
     def predict_peptides(self, peptides):
         """
@@ -114,6 +180,7 @@ def predict_peptides(self, peptides):
         df = self.predictor.affinity_predictor.predict_to_dataframe(
             peptides=batch_peptides,
             alleles=batch_alleles,
+            include_percentile_ranks=self.include_affinity_percentile_ranks,
         )
         binding_predictions = []
         for row in df.itertuples(index=False):
@@ -123,7 +190,7 @@ def predict_peptides(self, peptides):
                 affinity=row.prediction,
                 percentile_rank=(
                     row.prediction_percentile
-                    if hasattr(row, 'prediction_percentile') else nan),
+                    if hasattr(row, 'prediction_percentile') else None),
                 prediction_method_name="mhcflurry",
             ))
         return BindingPredictionCollection(binding_predictions)
@@ -152,6 +219,7 @@ def predict(self, peptides):
         aff_df = self.predictor.affinity_predictor.predict_to_dataframe(
             peptides=batch_peptides,
             alleles=batch_alleles,
+            include_percentile_ranks=self.include_affinity_percentile_ranks,
         )
 
         # Per-allele presentation calls (presentation predictor does
@@ -233,7 +301,8 @@ def __init__(
             alleles,
             default_peptide_lengths=[9],
             predictor=None,
-            models_path=None):
+            models_path=None,
+            include_affinity_percentile_ranks=True):
         """
         Parameters
         -----------
@@ -246,6 +315,12 @@ def __init__(
 
         models_path : string
             Models dir to use if predictor argument is None
+
+        include_affinity_percentile_ranks : bool
+            Whether to request affinity percentile ranks. Enabled by default.
+            If enabled, requested alleles must have MHCflurry affinity
+            percentile-rank calibration, either directly or through an allele
+            with the same pseudosequence.
         """
         from mhcflurry import Class1AffinityPredictor
         BasePredictor.__init__(
@@ -269,9 +344,14 @@ def __init__(
                         Class1AffinityPredictor.load()
             self.predictor = _model_cache[cache_key]
 
+        self.include_affinity_percentile_ranks = \
+            include_affinity_percentile_ranks
+
         for allele in self.alleles:
             if allele not in self.predictor.supported_alleles:
                 raise UnsupportedAllele(allele)
+        if self.include_affinity_percentile_ranks:
+            _check_affinity_percent_rank_support(self.predictor, self.alleles)
 
     def predict_peptides(self, peptides):
         """
@@ -286,6 +366,7 @@ def predict_peptides(self, peptides):
         df = self.predictor.predict_to_dataframe(
             peptides=batch_peptides,
             alleles=batch_alleles,
+            include_percentile_ranks=self.include_affinity_percentile_ranks,
         )
         binding_predictions = []
         for row in df.itertuples(index=False):
@@ -295,7 +376,7 @@ def predict_peptides(self, peptides):
                 affinity=row.prediction,
                 percentile_rank=(
                     row.prediction_percentile
-                    if hasattr(row, 'prediction_percentile') else nan),
+                    if hasattr(row, 'prediction_percentile') else None),
                 prediction_method_name="mhcflurry",
             ))
         return BindingPredictionCollection(binding_predictions)

tests/test_mhcflurry_key_lookup.py

@@ -17,21 +17,39 @@
 import pandas as pd
 import pytest
 
-from mhctools import MHCflurry
+from mhctools import MHCflurry, MHCflurry_Affinity
 
 
-def _make_fake_predictor(aff_allele_str, pres_allele_str, supported):
+def _make_fake_predictor(
+        aff_allele_str,
+        pres_allele_str,
+        supported,
+        percent_rank_transforms=None,
+        allele_to_sequence=None):
     """Build a fake mhcflurry Class1PresentationPredictor with configurable
     allele string in the affinity and presentation outputs."""
-    affinity_predictor = types.SimpleNamespace(
-        predict_to_dataframe=lambda peptides, alleles: pd.DataFrame({
+    def predict_to_dataframe(
+            peptides, alleles, include_percentile_ranks=True):
+        data = {
             "peptide": peptides,
             "allele": [aff_allele_str] * len(peptides),
             "prediction": [500.0] * len(peptides),
-            "prediction_percentile": [1.5] * len(peptides),
-        }),
+        }
+        if include_percentile_ranks:
+            data["prediction_percentile"] = [1.5] * len(peptides)
+        return pd.DataFrame(data)
+
+    affinity_predictor = types.SimpleNamespace(
+        predict_to_dataframe=predict_to_dataframe,
         supported_alleles=supported,
     )
+    if percent_rank_transforms is not None:
+        affinity_predictor.allele_to_percent_rank_transform = \
+            percent_rank_transforms
+    if allele_to_sequence is not None:
+        affinity_predictor.allele_to_sequence = allele_to_sequence
+        affinity_predictor.canonicalize_allele_name = lambda allele: allele
+
     def predict(peptides, alleles, include_affinity_percentile=False, verbose=0):
         return pd.DataFrame({
             "peptide": list(peptides),
@@ -70,3 +88,71 @@ def test_inconsistent_allele_strings_raise_instead_of_silently_returning_zero():
     p = MHCflurry(alleles=["HLA-A*02:01"], predictor=fake)
     with pytest.raises(ValueError, match="missing presentation score"):
         p.predict(["SIINFEKLA"])
+
+
+def test_accepts_affinity_percentile_calibration_from_same_pseudosequence():
+    fake = _make_fake_predictor(
+        aff_allele_str="HLA-C*15:05",
+        pres_allele_str="HLA-C*15:05",
+        supported=["HLA-C*15:05"],
+        percent_rank_transforms={"HLA-C*15:99": object()},
+        allele_to_sequence={
+            "HLA-C*15:05": "PSEUDOSEQ",
+            "HLA-C*15:99": "PSEUDOSEQ",
+        })
+    predictor = MHCflurry(alleles=["HLA-C*15:05"], predictor=fake)
+    results = predictor.predict(["SIINFEKLA"])
+    assert results[0].affinity.percentile_rank == 1.5
+
+
+def test_missing_affinity_percentile_calibration_raises_early():
+    fake = _make_fake_predictor(
+        aff_allele_str="HLA-C*15:05",
+        pres_allele_str="HLA-C*15:05",
+        supported=["HLA-C*15:05"],
+        percent_rank_transforms={},
+        allele_to_sequence={"HLA-C*15:05": "PSEUDOSEQ"})
+    with pytest.raises(ValueError, match="affinity percentile ranks"):
+        MHCflurry(alleles=["HLA-C*15:05"], predictor=fake)
+
+
+def test_can_disable_missing_affinity_percentile_ranks():
+    fake = _make_fake_predictor(
+        aff_allele_str="HLA-C*15:05",
+        pres_allele_str="HLA-C*15:05",
+        supported=["HLA-C*15:05"],
+        percent_rank_transforms={},
+        allele_to_sequence={"HLA-C*15:05": "PSEUDOSEQ"})
+    predictor = MHCflurry(
+        alleles=["HLA-C*15:05"],
+        predictor=fake,
+        include_affinity_percentile_ranks=False)
+    results = predictor.predict(["SIINFEKLA"])
+    assert results[0].affinity.percentile_rank is None
+
+
+def test_affinity_only_disabled_percentile_ranks_convert_to_none():
+    def predict_to_dataframe(
+            peptides, alleles, include_percentile_ranks=True):
+        data = {
+            "peptide": list(peptides),
+            "allele": list(alleles),
+            "prediction": [500.0] * len(peptides),
+        }
+        if include_percentile_ranks:
+            data["prediction_percentile"] = [1.5] * len(peptides)
+        return pd.DataFrame(data)
+
+    fake = types.SimpleNamespace(
+        predict_to_dataframe=predict_to_dataframe,
+        supported_alleles=["HLA-A*02:232"],
+    )
+    predictor = MHCflurry_Affinity(
+        alleles=["HLA-A*02:232"],
+        predictor=fake,
+        include_affinity_percentile_ranks=False)
+
+    result = predictor.predict(["SIINFEKLA"])[0]
+
+    assert result.affinity.percentile_rank is None
+    assert result.best_affinity_by_rank is None