diff --git a/README.md b/README.md index 8d2ee69..8c0947d 100644 --- a/README.md +++ b/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,49 @@ 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` | +| `NetMHCIIpan4_EL` | `pMHC_presentation` | `single_allele` | `II` | +| `NetMHCstabpan` | `pMHC_stability` | `single_allele` | `I` | +| `MHCflurry` | `pMHC_affinity` | `single_allele` | `I` | +| `MHCflurry` haplotype mode | `pMHC_presentation` | `haplotype` | `I` | +| `MHCflurry` per-allele panel mode | `pMHC_presentation` | `single_allele` | `I` | +| `Pepsickle` | `proteasome_cleavage` | `none` | `none` | + +For MHCflurry presentation, `presentation_allele_mode="haplotype"` treats the +requested alleles as one sample genotype and emits one `pMHC_presentation` +record per peptide. The `allele` field carries MHCflurry's `best_allele` +attribution when available. `presentation_allele_mode="per_allele"` treats each +allele as a separate one-allele synthetic sample and emits one presentation +record per peptide/allele pair. The default `"auto"` mode uses haplotype mode +for up to six alleles and per-allele mode for larger allele panels. + ### The Prediction object Every prediction is a frozen, self-contained `Prediction` dataclass: diff --git a/mhctools/__init__.py b/mhctools/__init__.py index d62cf65..29ef8ff 100644 --- a/mhctools/__init__.py +++ b/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,7 +72,7 @@ def __getattr__(name): raise AttributeError( "module %r has no attribute %r" % (__name__, name)) -__version__ = "3.13.6" +__version__ = "3.13.7" __all__ = [ "Prediction", @@ -71,6 +80,8 @@ def __getattr__(name): "PeptideResult", "PeptidePreds", # backward compat alias "Kind", + "MHC_CLASS_VALUES", + "MHC_DEPENDENCE_VALUES", "preds_from_rows", "MultiSample", "BindingPrediction", diff --git a/mhctools/base_predictor.py b/mhctools/base_predictor.py index b0af1bf..78c4889 100644 --- a/mhctools/base_predictor.py +++ b/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): diff --git a/mhctools/bigmhc.py b/mhctools/bigmhc.py index f5fb977..cdfe027 100644 --- a/mhctools/bigmhc.py +++ b/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 diff --git a/mhctools/iedb.py b/mhctools/iedb.py index a950114..98b93c2 100644 --- a/mhctools/iedb.py +++ b/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, diff --git a/mhctools/mhcflurry.py b/mhctools/mhcflurry.py index 07e2246..61d61d8 100644 --- a/mhctools/mhcflurry.py +++ b/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,14 +99,22 @@ 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. For presentation, + ``presentation_allele_mode`` controls whether mhctools treats the allele + set as one class-I haplotype or as a panel of independent one-allele + samples. The legacy ``predict_peptides`` method returns BindingPrediction + objects based on affinity values for backward compat. See https://github.com/openvax/mhcflurry """ uses_flanking_sequences = True flank_length = 15 + max_haplotype_alleles = 6 + presentation_allele_modes = frozenset(( + "auto", + "haplotype", + "per_allele", + )) def __init__( self, @@ -114,7 +122,8 @@ def __init__( default_peptide_lengths=[9], predictor=None, models_path=None, - include_affinity_percentile_ranks=True): + include_affinity_percentile_ranks=True, + presentation_allele_mode="auto"): """ Parameters ----------- @@ -133,6 +142,15 @@ def __init__( If enabled, requested alleles must have MHCflurry affinity percentile-rank calibration, either directly or through an allele with the same pseudosequence. + + presentation_allele_mode : {"auto", "haplotype", "per_allele"} + How to interpret the requested alleles for presentation scoring. + ``"haplotype"`` treats the alleles as one sample genotype and + emits one presentation prediction per peptide. ``"per_allele"`` + treats each allele as a separate one-allele synthetic sample and + emits one presentation prediction per peptide/allele pair. + ``"auto"`` uses haplotype mode for up to six alleles and + per-allele mode for larger allele panels. """ from mhcflurry import Class1PresentationPredictor BasePredictor.__init__( @@ -158,6 +176,8 @@ def __init__( self.include_affinity_percentile_ranks = \ include_affinity_percentile_ranks + self.presentation_allele_mode = self._resolve_presentation_allele_mode( + presentation_allele_mode) for allele in self.alleles: if allele not in self.predictor.supported_alleles: @@ -166,6 +186,27 @@ def __init__( _check_affinity_percent_rank_support( self.predictor.affinity_predictor, self.alleles) + def _resolve_presentation_allele_mode(self, presentation_allele_mode): + if presentation_allele_mode not in self.presentation_allele_modes: + raise ValueError( + "presentation_allele_mode must be one of %s, got %r" % ( + sorted(self.presentation_allele_modes), + presentation_allele_mode)) + if presentation_allele_mode == "auto": + if len(self.alleles) <= self.max_haplotype_alleles: + return "haplotype" + return "per_allele" + if ( + presentation_allele_mode == "haplotype" and + len(self.alleles) > self.max_haplotype_alleles): + raise ValueError( + "MHCflurry presentation haplotype mode accepts at most %d " + "alleles, got %d. Use presentation_allele_mode='per_allele' " + "for allele panels." % ( + self.max_haplotype_alleles, + len(self.alleles))) + return presentation_allele_mode + def _predict_protein_flank_lengths(self): processing_predictor = getattr( self.predictor, "processing_predictor_with_flanks", None) @@ -214,11 +255,14 @@ 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. Presentation predictions are haplotype-level + when ``presentation_allele_mode`` is ``"haplotype"`` and per-allele + when it is ``"per_allele"``. - Uses batch prediction across all alleles in a single call for - both affinity and presentation scores. + Uses batch prediction across alleles for affinity. For presentation, + haplotype mode passes the allele list as one MHCflurry sample genotype; + per-allele mode passes a sample-to-one-allele dict. """ from .pred import PeptideResult @@ -243,34 +287,58 @@ 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 self.presentation_allele_mode == "haplotype": + presentation_alleles = allele_list + else: + presentation_alleles = { + allele: [allele] + for allele in allele_list } - 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): + + kwargs = { + "peptides": peptide_list, + "alleles": presentation_alleles, + "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) + expected_presentation_rows = len(peptide_list) + if self.presentation_allele_mode == "per_allele": + expected_presentation_rows *= len(allele_list) + if len(pres_df) != expected_presentation_rows: + raise ValueError( + "MHCflurry returned %d presentation row(s) for %d " + "expected peptide/context row(s)" % ( + len(pres_df), + expected_presentation_rows)) + + pres_by_peptide_index = {i: [] for i in range(len(peptide_list))} + seen_presentation_keys = set() + for row_position, row in enumerate(pres_df.itertuples(index=False)): + row_index = int(getattr(row, "peptide_num", row_position)) + if self.presentation_allele_mode == "haplotype": + allele = getattr(row, "best_allele", getattr(row, "allele", "")) + key = (row_index, "") + else: + allele = getattr( + row, + "sample_name", + getattr(row, "best_allele", getattr(row, "allele", ""))) + key = (row_index, allele) + if key in seen_presentation_keys: 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 and allele '%s'" % (row_index, allele)) + seen_presentation_keys.add(key) + pres_by_peptide_index[row_index].append(( + row.presentation_score, + row.presentation_percentile, + allele, + )) groups = [list() for _ in peptide_list] for row_index, row in zip(batch_indices, aff_df.itertuples(index=False)): @@ -301,24 +369,27 @@ 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 not pres_by_peptide_index[row_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] - groups[row_index].append(Prediction( - kind=Kind.pMHC_presentation, - score=pres_score, - peptide=pep, - allele=allele, - n_flank=n_flank, - c_flank=c_flank, - percentile_rank=pres_pct, - predictor_name="mhcflurry", - )) + "peptide index %d, peptide='%s'" % (row_index, pep)) + 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 "") + for pres_score, pres_pct, presentation_allele in ( + pres_by_peptide_index[row_index]): + groups[row_index].append(Prediction( + kind=Kind.pMHC_presentation, + score=pres_score, + peptide=pep, + allele=presentation_allele or "", + n_flank=n_flank, + c_flank=c_flank, + percentile_rank=pres_pct, + predictor_name="mhcflurry", + )) return [PeptideResult(preds=tuple(preds)) for preds in groups] @@ -331,6 +402,22 @@ def predict_with_flanks(self, peptides, n_flanks, c_flanks): def _default_pred_kind(self): return Kind.pMHC_affinity + def kind_support(self): + presentation_dependence = ( + "haplotype" + if self.presentation_allele_mode == "haplotype" + else "single_allele") + return { + Kind.pMHC_affinity: { + "mhc_dependence": "single_allele", + "mhc_class": "I", + }, + Kind.pMHC_presentation: { + "mhc_dependence": presentation_dependence, + "mhc_class": "I", + }, + } + class MHCflurry_Affinity(BasePredictor): """ diff --git a/mhctools/netmhc_pan4.py b/mhctools/netmhc_pan4.py index d33fd69..af2b409 100644 --- a/mhctools/netmhc_pan4.py +++ b/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 diff --git a/mhctools/netmhc_pan41.py b/mhctools/netmhc_pan41.py index 16e6355..099ce8f 100644 --- a/mhctools/netmhc_pan41.py +++ b/mhctools/netmhc_pan41.py @@ -10,9 +10,11 @@ # 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_netmhc41_stdout, parse_netmhcpan_to_preds -from functools import partial +from .pred import Kind class NetMHCpan41(BaseCommandlinePredictor): @@ -38,6 +40,7 @@ def __init__( flags = [] else: raise ValueError("Unsupported mode", mode) + self.mode = mode BaseCommandlinePredictor.__init__( self, @@ -53,6 +56,20 @@ def __init__( extra_flags=flags + extra_flags, process_limit=process_limit) + def kind_support(self): + if self.mode == "binding_affinity": + kinds = (Kind.pMHC_affinity, Kind.pMHC_presentation) + else: + kinds = (Kind.pMHC_presentation,) + return { + kind: { + "mhc_dependence": "single_allele", + "mhc_class": "I", + } + for kind in kinds + } + + class NetMHCpan41_EL(NetMHCpan41): """ Wrapper for NetMHCpan4 when the preferred mode is elution score @@ -92,4 +109,4 @@ def __init__( program_name=program_name, process_limit=process_limit, mode="binding_affinity", - extra_flags=extra_flags) \ No newline at end of file + extra_flags=extra_flags) diff --git a/mhctools/netmhc_pan42.py b/mhctools/netmhc_pan42.py index 1425b7a..2852c17 100644 --- a/mhctools/netmhc_pan42.py +++ b/mhctools/netmhc_pan42.py @@ -10,9 +10,11 @@ # 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_netmhcpan41_stdout, parse_netmhcpan_to_preds -from functools import partial +from .pred import Kind class NetMHCpan42(BaseCommandlinePredictor): @@ -39,6 +41,7 @@ def __init__( flags = [] else: raise ValueError("Unsupported mode", mode) + self.mode = mode BaseCommandlinePredictor.__init__( self, @@ -54,6 +57,19 @@ def __init__( extra_flags=flags + extra_flags, process_limit=process_limit) + def kind_support(self): + if self.mode == "binding_affinity": + kinds = (Kind.pMHC_affinity, Kind.pMHC_presentation) + else: + kinds = (Kind.pMHC_presentation,) + return { + kind: { + "mhc_dependence": "single_allele", + "mhc_class": "I", + } + for kind in kinds + } + class NetMHCpan42_EL(NetMHCpan42): """NetMHCpan 4.2 in elution score mode.""" diff --git a/mhctools/netmhcii_pan.py b/mhctools/netmhcii_pan.py index 14fa07a..fb0c011 100644 --- a/mhctools/netmhcii_pan.py +++ b/mhctools/netmhcii_pan.py @@ -19,11 +19,14 @@ from .base_commandline_predictor import BaseCommandlinePredictor from .parsing import parse_netmhciipan_stdout, parse_netmhciipan4_stdout, parse_netmhciipan43_stdout +from .pred import Kind logger = logging.getLogger(__name__) class NetMHCIIpanBase(BaseCommandlinePredictor): + mhc_class = "II" + def __init__( self, alleles, @@ -132,6 +135,7 @@ def __init__( if mode not in ['binding_affinity', 'elution_score']: raise ValueError("Unsupported mode", mode) + self.mode = mode # Always include binding affinity data (-BA flag), though the main score and %rank will # still be EL-based. This gives us access to the BA-based score and %rank columns. @@ -145,6 +149,11 @@ def __init__( default_peptide_lengths=default_peptide_lengths, extra_flags=['-BA'] + extra_flags) + def _default_pred_kind(self): + if self.mode == "binding_affinity": + return Kind.pMHC_affinity + return Kind.pMHC_presentation + class NetMHCIIpan4_EL(NetMHCIIpan4): """ @@ -236,6 +245,7 @@ def __init__( if mode not in ['binding_affinity', 'elution_score']: raise ValueError("Unsupported mode", mode) + self.mode = mode # Always include binding affinity data (-BA flag), though the main score and %rank will # still be EL-based. This gives us access to the BA-based score and %rank columns. @@ -249,6 +259,11 @@ def __init__( default_peptide_lengths=default_peptide_lengths, extra_flags=['-BA'] + extra_flags) + def _default_pred_kind(self): + if self.mode == "binding_affinity": + return Kind.pMHC_affinity + return Kind.pMHC_presentation + class NetMHCIIpan43_EL(NetMHCIIpan43): """ Wrapper for NetMHCIIpan43 when the preferred mode is elution score diff --git a/mhctools/netmhcstabpan.py b/mhctools/netmhcstabpan.py index c2ed6b0..a557a3c 100644 --- a/mhctools/netmhcstabpan.py +++ b/mhctools/netmhcstabpan.py @@ -12,6 +12,7 @@ from .base_commandline_predictor import BaseCommandlinePredictor from .parsing import parse_netmhcstabpan +from .pred import Kind class NetMHCstabpan(BaseCommandlinePredictor): def __init__( @@ -42,4 +43,7 @@ def predict_peptides(self, peptides): if len(peptide_lengths) > 1: raise ValueError("All peptides must be the same length") return super().predict_peptides(peptides) + + def _default_pred_kind(self): + return Kind.pMHC_stability diff --git a/mhctools/pred.py b/mhctools/pred.py index 78127a1..36abfdc 100644 --- a/mhctools/pred.py +++ b/mhctools/pred.py @@ -18,11 +18,30 @@ import pandas as pd +MHC_DEPENDENCE_VALUES = frozenset(( + "none", + "single_allele", + "haplotype", +)) +"""Allowed ``kind_support()[kind]["mhc_dependence"]`` values.""" + + +MHC_CLASS_VALUES = frozenset(( + "none", + "I", + "II", + "both", +)) +"""Allowed ``kind_support()[kind]["mhc_class"]`` values.""" + + class Kind: """String constants for prediction kinds. You can use ``Kind.pMHC_affinity`` or just ``"pMHC_affinity"`` — - they're the same string. + they're the same string. These constants name what is measured, but + predictor instances define the MHC context required for their supported + kinds through ``kind_support()``. """ pMHC_affinity = "pMHC_affinity" pMHC_presentation = "pMHC_presentation" diff --git a/mhctools/processing_predictor.py b/mhctools/processing_predictor.py index a747d98..80f11aa 100644 --- a/mhctools/processing_predictor.py +++ b/mhctools/processing_predictor.py @@ -39,7 +39,11 @@ _normalize_sequence_dict, _peptide_contexts, ) -from .pred import Prediction, Kind, PeptideResult +from .pred import ( + Kind, + PeptideResult, + Prediction, +) # ------------------------------------------------------------------ @@ -263,6 +267,20 @@ def _pred_kind(self): """Kind value for Prediction objects. Override in subclasses.""" return Kind.antigen_processing + def kind_support(self): + """Predictor-specific MHC context for supported prediction kinds.""" + return { + self._pred_kind(): { + "mhc_dependence": "none", + "mhc_class": "none", + } + } + + @property + def supported_kinds(self): + """Prediction kind strings this predictor can emit.""" + return tuple(self.kind_support()) + def _predictor_name(self): return self.__class__.__name__.lower() diff --git a/tests/test_mhcflurry.py b/tests/test_mhcflurry.py index 88b48c9..f1945a5 100644 --- a/tests/test_mhcflurry.py +++ b/tests/test_mhcflurry.py @@ -157,7 +157,7 @@ def test_mhcflurry_affinity_only(): def test_mhcflurry_multiple_alleles(): - """MHCflurry with multiple alleles produces predictions for each allele.""" + """MHCflurry affinity is per-allele; presentation is haplotype-level.""" alleles = ["HLA-A*02:01", "HLA-B*07:02"] predictor = MHCflurry(alleles=alleles) results = predictor.predict(["SIINFEKL"]) @@ -165,10 +165,13 @@ def test_mhcflurry_multiple_alleles(): eq_(1, len(results), "Expected one PeptideResult") r = results[0] - # Should have 2 alleles x 2 kinds = 4 predictions total - eq_(4, len(r.preds), "Expected 4 predictions (2 alleles x 2 kinds)") + # Two per-allele affinity predictions plus one haplotype-level + # presentation prediction with best_allele attribution. + eq_(3, len(r.preds), "Expected 3 predictions") + eq_(2, len(r.filter(kind=Kind.pMHC_affinity))) + eq_(1, len(r.filter(kind=Kind.pMHC_presentation))) - # Both alleles should be present + # Both alleles should be present through affinity predictions. assert r.alleles == set(alleles) # Both kinds should be present diff --git a/tests/test_mhcflurry_key_lookup.py b/tests/test_mhcflurry_key_lookup.py index 008ae51..35fd20b 100644 --- a/tests/test_mhcflurry_key_lookup.py +++ b/tests/test_mhcflurry_key_lookup.py @@ -4,20 +4,20 @@ # # http://www.apache.org/licenses/LICENSE-2.0 -"""Tests for the MHCflurry wrapper's presentation-score lookup. - -Verifies that when the presentation predictor's output allele string differs -from the affinity predictor's output allele string (e.g. different -normalization), the wrapper fails loudly instead of silently returning -score=0.0 for all presentations. -""" +"""Tests for the MHCflurry wrapper's presentation and affinity bookkeeping.""" import types import pandas as pd import pytest -from mhctools import MHCflurry, MHCflurry_Affinity +from mhctools import ( + MHC_CLASS_VALUES, + MHC_DEPENDENCE_VALUES, + MHCflurry, + MHCflurry_Affinity, +) +from mhctools.pred import Kind def _make_fake_predictor( @@ -32,9 +32,12 @@ def _make_fake_predictor( def predict_to_dataframe( peptides, alleles, include_percentile_ranks=True): + output_alleles = ( + [aff_allele_str] * len(peptides) + if aff_allele_str is not None else list(alleles)) data = { "peptide": peptides, - "allele": [aff_allele_str] * len(peptides), + "allele": output_alleles, "prediction": [500.0] * len(peptides), } if include_percentile_ranks: @@ -63,18 +66,38 @@ def predict( throw=True, affinity_model_kwargs=None, processing_batch_size="auto"): + allele_arg = { + k: list(v) + for (k, v) in alleles.items() + } if isinstance(alleles, dict) else list(alleles) predict_calls.append({ "peptides": list(peptides), - "alleles": list(alleles), + "alleles": allele_arg, "n_flanks": None if n_flanks is None else list(n_flanks), "c_flanks": None if c_flanks is None else list(c_flanks), }) - return pd.DataFrame({ - "peptide": list(peptides), - "allele": [pres_allele_str] * len(peptides), - "presentation_score": [0.75] * len(peptides), - "presentation_percentile": [2.0] * len(peptides), - }) + rows = [] + if isinstance(alleles, dict): + for sample_name in alleles: + for i, peptide in enumerate(peptides): + rows.append({ + "peptide": peptide, + "peptide_num": i, + "sample_name": sample_name, + "best_allele": sample_name, + "presentation_score": 0.75, + "presentation_percentile": 2.0, + }) + else: + for i, peptide in enumerate(peptides): + rows.append({ + "peptide": peptide, + "peptide_num": i, + "best_allele": pres_allele_str, + "presentation_score": 0.75, + "presentation_percentile": 2.0, + }) + return pd.DataFrame(rows) return types.SimpleNamespace( affinity_predictor=affinity_predictor, predict=predict, @@ -96,17 +119,15 @@ def test_consistent_allele_strings_produce_correct_scores(): assert r.presentation.allele == "HLA-A*02:01" -def test_inconsistent_allele_strings_raise_instead_of_silently_returning_zero(): - """Regression: previously the wrapper silently returned (0.0, None) for - presentation when the allele string in aff_df didn't match the key in - pres_by_pep_allele. Now it raises.""" +def test_presentation_best_allele_is_attribution_not_affinity_lookup_key(): fake = _make_fake_predictor( aff_allele_str="HLA-A*02:01", # affinity output - pres_allele_str="HLA-A0201", # presentation output (different!) + pres_allele_str="HLA-A0201", # presentation attribution supported=["HLA-A*02:01"]) p = MHCflurry(alleles=["HLA-A*02:01"], predictor=fake) - with pytest.raises(ValueError, match="missing presentation score"): - p.predict(["SIINFEKLA"]) + result = p.predict(["SIINFEKLA"])[0] + assert result.presentation.allele == "HLA-A0201" + assert result.presentation.score == 0.75 def test_accepts_affinity_percentile_calibration_from_same_pseudosequence(): @@ -150,6 +171,23 @@ def test_can_disable_missing_affinity_percentile_ranks(): assert results[0].affinity.percentile_rank is None +def test_mhcflurry_kind_support_marks_presentation_as_haplotype_level(): + fake = _make_fake_predictor( + aff_allele_str="HLA-A*02:01", + pres_allele_str="HLA-A*02:01", + supported=["HLA-A*02:01"]) + predictor = MHCflurry(alleles=["HLA-A*02:01"], predictor=fake) + + support = predictor.kind_support() + + assert support[Kind.pMHC_affinity]["mhc_dependence"] == "single_allele" + assert support[Kind.pMHC_presentation]["mhc_dependence"] == "haplotype" + assert support[Kind.pMHC_presentation]["mhc_class"] == "I" + assert support[Kind.pMHC_presentation]["mhc_dependence"] in ( + MHC_DEPENDENCE_VALUES) + assert support[Kind.pMHC_presentation]["mhc_class"] in MHC_CLASS_VALUES + + def test_mhcflurry_forwards_flanks_to_presentation_predictor(): fake = _make_fake_predictor( aff_allele_str="HLA-A*02:01", @@ -171,6 +209,80 @@ def test_mhcflurry_forwards_flanks_to_presentation_predictor(): assert results[1].presentation.c_flank == "YY" +def test_mhcflurry_presentation_uses_full_haplotype_once(): + fake = _make_fake_predictor( + aff_allele_str=None, + pres_allele_str="HLA-B*07:02", + supported=["HLA-A*02:01", "HLA-B*07:02"]) + predictor = MHCflurry( + alleles=["HLA-A*02:01", "HLA-B*07:02"], + predictor=fake) + + result = predictor.predict(["SIINFEKLA"])[0] + + assert len(fake.predict_calls) == 1 + assert set(fake.predict_calls[0]["alleles"]) == { + "HLA-A*02:01", + "HLA-B*07:02", + } + assert len(result.filter(kind=Kind.pMHC_affinity)) == 2 + assert len(result.filter(kind=Kind.pMHC_presentation)) == 1 + assert result.presentation.allele == "HLA-B*07:02" + + +def test_mhcflurry_large_panel_uses_one_sample_per_allele(): + alleles = [ + "HLA-A*01:01", + "HLA-A*02:01", + "HLA-A*03:01", + "HLA-B*07:02", + "HLA-B*08:01", + "HLA-C*07:01", + "HLA-C*07:02", + ] + fake = _make_fake_predictor( + aff_allele_str=None, + pres_allele_str="unused", + supported=alleles) + predictor = MHCflurry(alleles=alleles, predictor=fake) + + result = predictor.predict(["SIINFEKLA"])[0] + + assert predictor.presentation_allele_mode == "per_allele" + assert predictor.kind_support()[Kind.pMHC_presentation][ + "mhc_dependence"] == "single_allele" + assert len(fake.predict_calls) == 1 + presentation_alleles = fake.predict_calls[0]["alleles"] + assert set(presentation_alleles) == set(alleles) + assert all(v == [k] for k, v in presentation_alleles.items()) + assert len(result.filter(kind=Kind.pMHC_affinity)) == len(alleles) + presentations = result.filter(kind=Kind.pMHC_presentation) + assert len(presentations) == len(alleles) + assert {p.allele for p in presentations} == set(alleles) + + +def test_mhcflurry_rejects_large_explicit_haplotype(): + alleles = [ + "HLA-A*01:01", + "HLA-A*02:01", + "HLA-A*03:01", + "HLA-B*07:02", + "HLA-B*08:01", + "HLA-C*07:01", + "HLA-C*07:02", + ] + fake = _make_fake_predictor( + aff_allele_str=None, + pres_allele_str="unused", + supported=alleles) + + with pytest.raises(ValueError, match="haplotype mode accepts at most"): + MHCflurry( + alleles=alleles, + predictor=fake, + presentation_allele_mode="haplotype") + + def test_mhcflurry_predict_proteins_threads_flanking_context(): fake = _make_fake_predictor( aff_allele_str="HLA-A*02:01", diff --git a/tests/test_pepsickle.py b/tests/test_pepsickle.py index 08f094f..3ee7d0e 100644 --- a/tests/test_pepsickle.py +++ b/tests/test_pepsickle.py @@ -22,7 +22,7 @@ score_nterm_cterm_anti_max_internal, ) from mhctools.proteasome_predictor import ProteasomePredictor -from mhctools.pred import Kind, PeptideResult, COLUMNS +from mhctools.pred import COLUMNS, Kind, PeptideResult pepsickle = pytest.importorskip("pepsickle") @@ -67,6 +67,14 @@ def test_str(): assert "Pepsickle" in s +def test_kind_support_is_mhc_independent(): + support = Pepsickle().kind_support() + + assert set(support) == {Kind.proteasome_cleavage} + assert support[Kind.proteasome_cleavage]["mhc_dependence"] == "none" + assert support[Kind.proteasome_cleavage]["mhc_class"] == "none" + + # -- cleavage_probs -- def test_cleavage_probs(predictor): @@ -245,4 +253,3 @@ def test_scoring_methods_produce_different_scores(): unique = set( tuple(round(s, 6) for s in v) for v in scores_by_fn.values()) assert len(unique) > 1 - diff --git a/tests/test_pred.py b/tests/test_pred.py index d1f0870..9d486a0 100644 --- a/tests/test_pred.py +++ b/tests/test_pred.py @@ -10,11 +10,22 @@ # See the License for the specific language governing permissions and # limitations under the License. -from mhctools.pred import Prediction, PeptideResult, Kind, preds_from_rows, COLUMNS +from mhctools.pred import ( + COLUMNS, + Kind, + MHC_CLASS_VALUES, + MHC_DEPENDENCE_VALUES, + PeptideResult, + Prediction, + preds_from_rows, +) from mhctools.sample import MultiSample from mhctools.base_predictor import BasePredictor from mhctools.binding_prediction import BindingPrediction from mhctools.binding_prediction_collection import BindingPredictionCollection +from mhctools.netmhc_pan41 import NetMHCpan41 +from mhctools.netmhcstabpan import NetMHCstabpan +from mhctools.netmhcii_pan import NetMHCIIpan4, NetMHCIIpan43 from mhctools.random_predictor import RandomBindingPredictor @@ -347,6 +358,63 @@ def test_predict_proteins(): assert pred.source_sequence_name == "TP53" +def test_binding_predictor_kind_support_defaults_to_single_allele_class_i(): + predictor = RandomBindingPredictor( + alleles=["HLA-A*02:01"], + default_peptide_lengths=[9]) + support = predictor.kind_support() + assert predictor.supported_kinds == (Kind.pMHC_affinity,) + assert support[Kind.pMHC_affinity]["mhc_dependence"] == "single_allele" + assert support[Kind.pMHC_affinity]["mhc_class"] == "I" + assert support[Kind.pMHC_affinity]["mhc_dependence"] in ( + MHC_DEPENDENCE_VALUES) + assert support[Kind.pMHC_affinity]["mhc_class"] in MHC_CLASS_VALUES + + +def test_netmhcpan41_kind_support_includes_affinity_and_presentation(): + predictor = NetMHCpan41.__new__(NetMHCpan41) + predictor.mode = "binding_affinity" + + support = predictor.kind_support() + + assert set(support) == {Kind.pMHC_affinity, Kind.pMHC_presentation} + assert support[Kind.pMHC_affinity]["mhc_dependence"] == "single_allele" + assert support[Kind.pMHC_presentation]["mhc_dependence"] == "single_allele" + assert support[Kind.pMHC_affinity]["mhc_class"] == "I" + + +def test_netmhciipan4_el_kind_support_is_class_ii_presentation(): + predictor = NetMHCIIpan4.__new__(NetMHCIIpan4) + predictor.mode = "elution_score" + + support = predictor.kind_support() + + assert set(support) == {Kind.pMHC_presentation} + assert support[Kind.pMHC_presentation]["mhc_dependence"] == "single_allele" + assert support[Kind.pMHC_presentation]["mhc_class"] == "II" + + +def test_netmhciipan43_ba_kind_support_is_class_ii_affinity(): + predictor = NetMHCIIpan43.__new__(NetMHCIIpan43) + predictor.mode = "binding_affinity" + + support = predictor.kind_support() + + assert set(support) == {Kind.pMHC_affinity} + assert support[Kind.pMHC_affinity]["mhc_dependence"] == "single_allele" + assert support[Kind.pMHC_affinity]["mhc_class"] == "II" + + +def test_netmhcstabpan_kind_support_is_stability(): + predictor = NetMHCstabpan.__new__(NetMHCstabpan) + + support = predictor.kind_support() + + assert set(support) == {Kind.pMHC_stability} + assert support[Kind.pMHC_stability]["mhc_dependence"] == "single_allele" + assert support[Kind.pMHC_stability]["mhc_class"] == "I" + + class FlankEchoPredictor(BasePredictor): uses_flanking_sequences = True flank_length = 2