diff --git a/mhctools/__init__.py b/mhctools/__init__.py index df5ab32..c4063ec 100644 --- a/mhctools/__init__.py +++ b/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", diff --git a/mhctools/mhcflurry.py b/mhctools/mhcflurry.py index 6583114..0c91a57 100644 --- a/mhctools/mhcflurry.py +++ b/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) diff --git a/tests/test_mhcflurry_key_lookup.py b/tests/test_mhcflurry_key_lookup.py index 884aa75..704f083 100644 --- a/tests/test_mhcflurry_key_lookup.py +++ b/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