From 5f393d7d5a26239b246e0003d5a338745dfbd911 Mon Sep 17 00:00:00 2001 From: Alex Rubinsteyn Date: Wed, 6 May 2026 10:53:10 -0400 Subject: [PATCH 1/2] Add flank-aware protein scanning --- mhctools/__init__.py | 2 +- mhctools/base_commandline_predictor.py | 7 +- mhctools/base_predictor.py | 237 +++++++++++++++++++------ mhctools/mhcflurry.py | 80 +++++++-- mhctools/processing_predictor.py | 61 ++++--- tests/test_mhcflurry.py | 36 ++++ tests/test_mhcflurry_key_lookup.py | 64 ++++++- tests/test_pred.py | 56 ++++++ 8 files changed, 440 insertions(+), 103 deletions(-) diff --git a/mhctools/__init__.py b/mhctools/__init__.py index cda4714..d62cf65 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.5" +__version__ = "3.13.6" __all__ = [ "Prediction", diff --git a/mhctools/base_commandline_predictor.py b/mhctools/base_commandline_predictor.py index 39f8e8a..b562d96 100644 --- a/mhctools/base_commandline_predictor.py +++ b/mhctools/base_commandline_predictor.py @@ -326,7 +326,7 @@ def _run_commands_and_collect_preds( groups[key].append(pred) return [PeptideResult(preds=tuple(preds)) for preds in groups.values()] - def predict(self, peptides): + def predict(self, peptides, n_flanks=None, c_flanks=None): """ Predict for a list of peptide sequences. @@ -334,8 +334,11 @@ def predict(self, peptides): available, parses directly to Pred objects. Otherwise falls back to converting from BindingPrediction. """ + peptides, n_flank_list, c_flank_list = self._check_flank_inputs( + peptides, n_flanks, c_flanks) if self.parse_to_preds_fn is None: - return super().predict(peptides) + return super().predict( + peptides, n_flanks=n_flank_list, c_flanks=c_flank_list) self._check_peptide_inputs(peptides) input_filenames = create_input_peptides_files( diff --git a/mhctools/base_predictor.py b/mhctools/base_predictor.py index 4c581d3..b0af1bf 100644 --- a/mhctools/base_predictor.py +++ b/mhctools/base_predictor.py @@ -12,7 +12,7 @@ import logging import warnings -from collections import defaultdict +from collections import defaultdict, namedtuple from typechecks import require_iterable_of from .allele_normalization import normalize_allele_name @@ -23,10 +23,85 @@ logger = logging.getLogger(__name__) +PeptideContext = namedtuple( + "PeptideContext", + ["source_sequence_name", "offset", "peptide", "n_flank", "c_flank"]) + + +def _normalize_sequence_dict(sequence_dict): + if isinstance(sequence_dict, str): + return {"seq": sequence_dict} + if isinstance(sequence_dict, (list, tuple)): + return {seq: seq for seq in sequence_dict} + return sequence_dict + + +def _check_flank_inputs(peptides, n_flanks=None, c_flanks=None): + peptide_list = list(peptides) + + def _check_flanks(name, flanks): + if flanks is None: + return None + if isinstance(flanks, str): + raise TypeError("%s must be a sequence of strings, not a string" % name) + flank_list = list(flanks) + require_iterable_of(flank_list, str, name) + if len(flank_list) != len(peptide_list): + raise ValueError( + "%s must have one entry per peptide, got %d flank(s) for " + "%d peptide(s)" % (name, len(flank_list), len(peptide_list))) + return flank_list + + return ( + peptide_list, + _check_flanks("n_flanks", n_flanks), + _check_flanks("c_flanks", c_flanks), + ) + + +def _peptide_contexts( + sequence_dict, + peptide_lengths, + flank_length=0, + n_flank_length=None, + c_flank_length=None): + if n_flank_length is None: + n_flank_length = flank_length + if c_flank_length is None: + c_flank_length = flank_length + + contexts = [] + for name, sequence in sequence_dict.items(): + for peptide_length in peptide_lengths: + for i in range(len(sequence) - peptide_length + 1): + peptide = sequence[i:i + peptide_length] + if n_flank_length or c_flank_length: + n_flank = sequence[max(0, i - n_flank_length):i] + c_flank = sequence[ + i + peptide_length: + i + peptide_length + c_flank_length] + else: + n_flank = "" + c_flank = "" + contexts.append(PeptideContext( + source_sequence_name=name, + offset=i, + peptide=peptide, + n_flank=n_flank, + c_flank=c_flank, + )) + return contexts + + class BasePredictor(object): """ Base class for all MHC binding predictors. """ + uses_flanking_sequences = False + flank_length = 15 + n_flank_length = None + c_flank_length = None + def __init__( self, alleles, @@ -90,26 +165,85 @@ def __str__(self): # --- new API --- - def predict(self, peptides): + def predict(self, peptides, n_flanks=None, c_flanks=None): """ Predict for a list of peptide sequences. + n_flanks and c_flanks are accepted for a uniform flank-aware API. + The default BindingPrediction-based implementation validates but + ignores them; subclasses that use flanking context should override + this method and set ``uses_flanking_sequences = True``. + Returns ------- list of PeptideResult """ - collection = self.predict_peptides(peptides) + peptide_list, _, _ = _check_flank_inputs( + peptides, n_flanks, c_flanks) + collection = self.predict_peptides(peptide_list) return collection.to_peptide_preds(kind=self._default_pred_kind()) - def predict_dataframe(self, peptides, sample_name=""): + def predict_with_flanks(self, peptides, n_flanks, c_flanks): + """ + Optional flank-aware prediction path. + + The default implementation validates aligned flank lists and falls + back to ``predict(peptides)`` so predictors that do not use flanking + context retain their existing behavior. Subclasses whose models use + flanks should override this method and forward the flanks to the + underlying predictor. + """ + peptide_list, _, _ = _check_flank_inputs(peptides, n_flanks, c_flanks) + return self.predict(peptide_list) + + def predict_dataframe( + self, peptides, sample_name="", n_flanks=None, c_flanks=None): """predict() flattened to a DataFrame.""" import pandas as pd - dfs = [pp.to_dataframe(sample_name) for pp in self.predict(peptides)] + dfs = [ + pp.to_dataframe(sample_name) + for pp in self.predict( + peptides, n_flanks=n_flanks, c_flanks=c_flanks) + ] if not dfs: from .pred import COLUMNS return pd.DataFrame(columns=COLUMNS) return pd.concat(dfs, ignore_index=True) + def _predict_protein_flank_length(self): + return max(self._predict_protein_flank_lengths()) + + def _predict_protein_flank_lengths(self): + if not self.uses_flanking_sequences: + return (0, 0) + n_flank_length = ( + self.flank_length + if self.n_flank_length is None else self.n_flank_length) + c_flank_length = ( + self.flank_length + if self.c_flank_length is None else self.c_flank_length) + return (n_flank_length, c_flank_length) + + def _check_flank_inputs(self, peptides, n_flanks=None, c_flanks=None): + return _check_flank_inputs(peptides, n_flanks, c_flanks) + + @staticmethod + def _with_protein_location(pred, context): + return Prediction( + kind=pred.kind, + score=pred.score, + peptide=pred.peptide, + allele=pred.allele, + n_flank=context.n_flank, + c_flank=context.c_flank, + value=pred.value, + percentile_rank=pred.percentile_rank, + source_sequence_name=context.source_sequence_name, + offset=context.offset, + predictor_name=pred.predictor_name, + predictor_version=pred.predictor_version, + ) + def predict_proteins(self, sequence_dict, peptide_lengths=None): """ Scan protein sequences and predict for all subsequences. @@ -126,23 +260,44 @@ def predict_proteins(self, sequence_dict, peptide_lengths=None): ------- dict mapping sequence_name -> list of PeptideResult """ - if isinstance(sequence_dict, str): - sequence_dict = {"seq": sequence_dict} - elif isinstance(sequence_dict, (list, tuple)): - sequence_dict = {seq: seq for seq in sequence_dict} + sequence_dict = _normalize_sequence_dict(sequence_dict) peptide_lengths = self._check_peptide_lengths(peptide_lengths) - peptide_set = set() - peptide_to_name_offset_pairs = defaultdict(list) - - for name, sequence in sequence_dict.items(): - for peptide_length in peptide_lengths: - for i in range(len(sequence) - peptide_length + 1): - peptide = sequence[i:i + peptide_length] - peptide_set.add(peptide) - peptide_to_name_offset_pairs[peptide].append((name, i)) + n_flank_length, c_flank_length = self._predict_protein_flank_lengths() + contexts = _peptide_contexts( + sequence_dict, + peptide_lengths, + n_flank_length=n_flank_length, + c_flank_length=c_flank_length) + + if self.uses_flanking_sequences: + peptide_list = [context.peptide for context in contexts] + n_flanks = [context.n_flank for context in contexts] + c_flanks = [context.c_flank for context in contexts] + flat_preds = self.predict_with_flanks( + peptide_list, + n_flanks=n_flanks, + c_flanks=c_flanks) + if len(flat_preds) != len(contexts): + raise ValueError( + "%s.predict returned %d result(s) for %d flanked " + "peptide occurrence(s)" % ( + self.__class__.__name__, len(flat_preds), + len(contexts))) + results = defaultdict(list) + for context, pp in zip(contexts, flat_preds): + relocated = PeptideResult(preds=tuple( + self._with_protein_location(p, context) + for p in pp.preds + )) + results[context.source_sequence_name].append(relocated) + return dict(results) + peptide_set = {context.peptide for context in contexts} + peptide_to_contexts = defaultdict(list) + for context in contexts: + peptide_to_contexts[context.peptide].append(context) peptide_list = sorted(peptide_set) flat_preds = self.predict(peptide_list) @@ -152,22 +307,12 @@ def predict_proteins(self, sequence_dict, peptide_lengths=None): if not pp.preds: continue peptide = pp.preds[0].peptide - for name, offset in peptide_to_name_offset_pairs.get(peptide, []): + for context in peptide_to_contexts.get(peptide, []): relocated = PeptideResult(preds=tuple( - Prediction( - kind=p.kind, - score=p.score, - peptide=p.peptide, - allele=p.allele, - value=p.value, - percentile_rank=p.percentile_rank, - source_sequence_name=name, - offset=offset, - predictor_name=p.predictor_name, - predictor_version=p.predictor_version, - ) for p in pp.preds + self._with_protein_location(p, context) + for p in pp.preds )) - results[name].append(relocated) + results[context.source_sequence_name].append(relocated) return dict(results) def predict_proteins_dataframe(self, sequence_dict, peptide_lengths=None, sample_name=""): @@ -312,25 +457,15 @@ def predict_subsequences( and an optional list of peptide lengths, returns a BindingPredictionCollection. """ - if isinstance(sequence_dict, str): - sequence_dict = {"seq": sequence_dict} - elif isinstance(sequence_dict, (list, tuple)): - sequence_dict = {seq: seq for seq in sequence_dict} + sequence_dict = _normalize_sequence_dict(sequence_dict) peptide_lengths = self._check_peptide_lengths(peptide_lengths) - # convert long protein sequences to set of peptides and - # associated sequence name / offsets that each peptide may have come - # from - peptide_set = set([]) - peptide_to_name_offset_pairs = defaultdict(list) - - for name, sequence in sequence_dict.items(): - for peptide_length in peptide_lengths: - for i in range(len(sequence) - peptide_length + 1): - peptide = sequence[i:i + peptide_length] - peptide_set.add(peptide) - peptide_to_name_offset_pairs[peptide].append((name, i)) + contexts = _peptide_contexts(sequence_dict, peptide_lengths) + peptide_set = {context.peptide for context in contexts} + peptide_to_contexts = defaultdict(list) + for context in contexts: + peptide_to_contexts[context.peptide].append(context) peptide_list = sorted(peptide_set) binding_predictions = self.predict_peptides(peptide_list) @@ -339,10 +474,10 @@ def predict_subsequences( results = [] for binding_prediction in binding_predictions: peptide = binding_prediction.peptide - for name, offset in peptide_to_name_offset_pairs[peptide]: + for context in peptide_to_contexts[peptide]: results.append(binding_prediction.clone_with_updates( - source_sequence_name=name, - offset=offset)) + source_sequence_name=context.source_sequence_name, + offset=context.offset)) self._check_results( results, peptides=peptide_set, diff --git a/mhctools/mhcflurry.py b/mhctools/mhcflurry.py index 0c91a57..07e2246 100644 --- a/mhctools/mhcflurry.py +++ b/mhctools/mhcflurry.py @@ -15,6 +15,7 @@ import os from .base_predictor import BasePredictor +from .base_predictor import _check_flank_inputs from .binding_prediction import BindingPrediction from .binding_prediction_collection import BindingPredictionCollection from .pred import Prediction, Kind @@ -104,6 +105,8 @@ class MHCflurry(BasePredictor): See https://github.com/openvax/mhcflurry """ + uses_flanking_sequences = True + flank_length = 15 def __init__( self, @@ -163,6 +166,18 @@ def __init__( _check_affinity_percent_rank_support( self.predictor.affinity_predictor, self.alleles) + def _predict_protein_flank_lengths(self): + processing_predictor = getattr( + self.predictor, "processing_predictor_with_flanks", None) + sequence_lengths = getattr( + processing_predictor, "sequence_lengths", None) + if sequence_lengths: + return ( + int(sequence_lengths.get("n_flank", self.flank_length)), + int(sequence_lengths.get("c_flank", self.flank_length)), + ) + return super()._predict_protein_flank_lengths() + def predict_peptides(self, peptides): """ Predict MHC binding affinity and presentation for peptides. @@ -195,7 +210,7 @@ def predict_peptides(self, peptides): )) return BindingPredictionCollection(binding_predictions) - def predict(self, peptides): + def predict(self, peptides, n_flanks=None, c_flanks=None): """ Predict for a list of peptide sequences. @@ -205,15 +220,21 @@ def predict(self, peptides): Uses batch prediction across all alleles in a single call for both affinity and presentation scores. """ - from collections import defaultdict from .pred import PeptideResult - peptide_list = list(peptides) + peptide_list, n_flank_list, c_flank_list = _check_flank_inputs( + peptides, n_flanks, c_flanks) + if n_flank_list is not None or c_flank_list is not None: + if n_flank_list is None: + n_flank_list = [""] * len(peptide_list) + if c_flank_list is None: + c_flank_list = [""] * len(peptide_list) allele_list = list(self.alleles) # Build cross product batch_peptides = peptide_list * len(allele_list) batch_alleles = [a for a in allele_list for _ in peptide_list] + batch_indices = list(range(len(peptide_list))) * len(allele_list) # Single batched call for affinity aff_df = self.predictor.affinity_predictor.predict_to_dataframe( @@ -228,21 +249,31 @@ def predict(self, peptides): # allele string so lookups with aff_df.allele always match. pres_by_pep_allele = {} for input_allele in allele_list: - df = self.predictor.predict( - peptides=peptide_list, - alleles=[input_allele], - include_affinity_percentile=False, - verbose=0, - ) - for row in df.itertuples(index=False): + 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): + 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.peptide, output_allele)] = ( + pres_by_pep_allele[(row_index, output_allele)] = ( row.presentation_score, row.presentation_percentile, ) - groups = defaultdict(list) - for row in aff_df.itertuples(index=False): + groups = [list() for _ in peptide_list] + for row_index, row in zip(batch_indices, aff_df.itertuples(index=False)): pep = row.peptide allele = row.allele affinity_nM = row.prediction @@ -253,17 +284,24 @@ def predict(self, peptides): aff_score = max(0.0, min(1.0, 1.0 - math.log(max(affinity_nM, 1e-6)) / math.log(50000))) - groups[pep].append(Prediction( + 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_affinity, score=aff_score, peptide=pep, allele=allele, + n_flank=n_flank, + c_flank=c_flank, value=affinity_nM, percentile_rank=affinity_pct, predictor_name="mhcflurry", )) - key = (pep, allele) + key = (row_index, allele) if key not in pres_by_pep_allele: raise ValueError( "MHCflurry: missing presentation score for " @@ -271,16 +309,24 @@ def predict(self, peptides): "peptide string mismatch between the affinity and " "presentation predictor outputs)" % (pep, allele)) pres_score, pres_pct = pres_by_pep_allele[key] - groups[pep].append(Prediction( + 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", )) - return [PeptideResult(preds=tuple(preds)) for preds in groups.values()] + return [PeptideResult(preds=tuple(preds)) for preds in groups] + + def predict_with_flanks(self, peptides, n_flanks, c_flanks): + return self.predict( + peptides, + n_flanks=n_flanks, + c_flanks=c_flanks) def _default_pred_kind(self): return Kind.pMHC_affinity diff --git a/mhctools/processing_predictor.py b/mhctools/processing_predictor.py index 7ffb0b7..a747d98 100644 --- a/mhctools/processing_predictor.py +++ b/mhctools/processing_predictor.py @@ -34,6 +34,11 @@ from collections import defaultdict +from .base_predictor import ( + _check_flank_inputs, + _normalize_sequence_dict, + _peptide_contexts, +) from .pred import Prediction, Kind, PeptideResult @@ -285,10 +290,13 @@ def predict(self, peptides, n_flanks=None, c_flanks=None): ------- list of PeptideResult """ + peptide_list, n_flank_list, c_flank_list = _check_flank_inputs( + peptides, n_flanks, c_flanks) + prediction_inputs = [] - for i, peptide in enumerate(peptides): - n_flank = n_flanks[i] if n_flanks else "" - c_flank = c_flanks[i] if c_flanks else "" + for i, peptide in enumerate(peptide_list): + n_flank = n_flank_list[i] if n_flank_list is not None else "" + c_flank = c_flank_list[i] if c_flank_list is not None else "" full_seq = n_flank + peptide + c_flank prediction_inputs.append((peptide, n_flank, c_flank, full_seq)) full_sequences = [item[3] for item in prediction_inputs] @@ -347,38 +355,30 @@ def predict_proteins(self, sequence_dict, peptide_lengths=None, ------- dict mapping sequence_name -> list of PeptideResult """ - if isinstance(sequence_dict, str): - sequence_dict = {"seq": sequence_dict} - elif isinstance(sequence_dict, (list, tuple)): - sequence_dict = {seq: seq for seq in sequence_dict} + sequence_dict = _normalize_sequence_dict(sequence_dict) peptide_lengths = self._resolve_peptide_lengths(peptide_lengths) probs_by_sequence = self.cleavage_probs_many(sequence_dict.values()) results = defaultdict(list) - for name, sequence in sequence_dict.items(): + for context in _peptide_contexts( + sequence_dict, peptide_lengths, flank_length): + sequence = sequence_dict[context.source_sequence_name] probs = probs_by_sequence[sequence] - for plen in peptide_lengths: - for i in range(len(sequence) - plen + 1): - peptide = sequence[i:i + plen] - score = self._peptide_score(probs, offset=i, length=plen) - if flank_length: - n_flank = sequence[max(0, i - flank_length):i] - c_flank = sequence[i + plen:i + plen + flank_length] - else: - n_flank = "" - c_flank = "" - pred = Prediction( - kind=self._pred_kind(), - score=score, - peptide=peptide, - n_flank=n_flank, - c_flank=c_flank, - source_sequence_name=name, - offset=i, - predictor_name=self._predictor_name(), - ) - results[name].append(PeptideResult(preds=(pred,))) + score = self._peptide_score( + probs, offset=context.offset, length=len(context.peptide)) + pred = Prediction( + kind=self._pred_kind(), + score=score, + peptide=context.peptide, + n_flank=context.n_flank, + c_flank=context.c_flank, + source_sequence_name=context.source_sequence_name, + offset=context.offset, + predictor_name=self._predictor_name(), + ) + results[context.source_sequence_name].append( + PeptideResult(preds=(pred,))) return dict(results) def predict_proteins_dataframe( @@ -404,8 +404,7 @@ def predict_cleavage_sites(self, sequence_dict): ------- dict mapping sequence_name -> list of float """ - if isinstance(sequence_dict, str): - sequence_dict = {"seq": sequence_dict} + sequence_dict = _normalize_sequence_dict(sequence_dict) probs_by_sequence = self.cleavage_probs_many(sequence_dict.values()) return { name: probs_by_sequence[seq] diff --git a/tests/test_mhcflurry.py b/tests/test_mhcflurry.py index 9ccf5bf..0e772d4 100644 --- a/tests/test_mhcflurry.py +++ b/tests/test_mhcflurry.py @@ -79,6 +79,42 @@ def test_mhcflurry_presentation_predict(): assert Kind.pMHC_presentation in r.kinds +def test_mhcflurry_predict_proteins_matches_direct_flanked_prediction(): + """Protein scanning should forward the same flanks as direct MHCflurry.""" + predictor = MHCflurry(alleles=[DEFAULT_ALLELE]) + peptide = "SIINFEKL" + protein = "MDSKG%sGSRLL" % peptide + offset = protein.index(peptide) + n_flank_length, c_flank_length = predictor._predict_protein_flank_lengths() + n_flank = protein[max(0, offset - n_flank_length):offset] + c_flank = protein[offset + len(peptide): + offset + len(peptide) + c_flank_length] + + protein_results = predictor.predict_proteins( + {"protein": protein}, + peptide_lengths=[len(peptide)]) + protein_result = [ + pp for pp in protein_results["protein"] if pp.offset == offset + ][0] + + direct = predictor.predictor.predict( + peptides=[peptide], + alleles=[DEFAULT_ALLELE], + n_flanks=[n_flank], + c_flanks=[c_flank], + include_affinity_percentile=False, + verbose=0) + + assert protein_result.presentation.n_flank == n_flank + assert protein_result.presentation.c_flank == c_flank + testing.assert_allclose( + protein_result.presentation.score, + direct.presentation_score.iloc[0]) + testing.assert_allclose( + protein_result.presentation.percentile_rank, + direct.presentation_percentile.iloc[0]) + + def test_mhcflurry_presentation_affinity_matches_old_api(): """Affinity values from the presentation predictor should be close to the old Class1AffinityPredictor values.""" diff --git a/tests/test_mhcflurry_key_lookup.py b/tests/test_mhcflurry_key_lookup.py index 704f083..008ae51 100644 --- a/tests/test_mhcflurry_key_lookup.py +++ b/tests/test_mhcflurry_key_lookup.py @@ -28,6 +28,8 @@ def _make_fake_predictor( allele_to_sequence=None): """Build a fake mhcflurry Class1PresentationPredictor with configurable allele string in the affinity and presentation outputs.""" + predict_calls = [] + def predict_to_dataframe( peptides, alleles, include_percentile_ranks=True): data = { @@ -50,7 +52,23 @@ def predict_to_dataframe( 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): + def predict( + peptides, + alleles, + sample_names=None, + n_flanks=None, + c_flanks=None, + include_affinity_percentile=False, + verbose=0, + throw=True, + affinity_model_kwargs=None, + processing_batch_size="auto"): + predict_calls.append({ + "peptides": list(peptides), + "alleles": list(alleles), + "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), @@ -60,6 +78,7 @@ def predict(peptides, alleles, include_affinity_percentile=False, verbose=0): return types.SimpleNamespace( affinity_predictor=affinity_predictor, predict=predict, + predict_calls=predict_calls, supported_alleles=supported, ) @@ -131,6 +150,49 @@ def test_can_disable_missing_affinity_percentile_ranks(): assert results[0].affinity.percentile_rank is None +def test_mhcflurry_forwards_flanks_to_presentation_predictor(): + 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) + + results = predictor.predict( + ["SIINFEKLA", "SIINFEKLA"], + n_flanks=["NN", "XX"], + c_flanks=["CC", "YY"]) + + assert fake.predict_calls[0]["n_flanks"] == ["NN", "XX"] + assert fake.predict_calls[0]["c_flanks"] == ["CC", "YY"] + assert len(results) == 2 + assert results[0].presentation.n_flank == "NN" + assert results[0].presentation.c_flank == "CC" + assert results[1].presentation.n_flank == "XX" + assert results[1].presentation.c_flank == "YY" + + +def test_mhcflurry_predict_proteins_threads_flanking_context(): + 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) + + result = predictor.predict_proteins({"protein": "MSIINFEKLAC"}) + + assert fake.predict_calls[0]["peptides"] == [ + "MSIINFEKL", + "SIINFEKLA", + "IINFEKLAC", + ] + assert fake.predict_calls[0]["n_flanks"] == ["", "M", "MS"] + assert fake.predict_calls[0]["c_flanks"] == ["AC", "C", ""] + middle = result["protein"][1] + assert middle.peptide == "SIINFEKLA" + assert middle.presentation.n_flank == "M" + assert middle.presentation.c_flank == "C" + + def test_affinity_only_disabled_percentile_ranks_convert_to_none(): def predict_to_dataframe( peptides, alleles, include_percentile_ranks=True): diff --git a/tests/test_pred.py b/tests/test_pred.py index 4d5e0b3..d1f0870 100644 --- a/tests/test_pred.py +++ b/tests/test_pred.py @@ -12,6 +12,7 @@ from mhctools.pred import Prediction, PeptideResult, Kind, preds_from_rows, COLUMNS 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.random_predictor import RandomBindingPredictor @@ -346,6 +347,61 @@ def test_predict_proteins(): assert pred.source_sequence_name == "TP53" +class FlankEchoPredictor(BasePredictor): + uses_flanking_sequences = True + flank_length = 2 + + def __init__(self): + BasePredictor.__init__( + self, + alleles=["HLA-A*02:01"], + default_peptide_lengths=[3], + min_peptide_length=3) + self.calls = [] + + def predict(self, peptides, n_flanks=None, c_flanks=None): + peptide_list, n_flank_list, c_flank_list = self._check_flank_inputs( + peptides, n_flanks, c_flanks) + self.calls.append((peptide_list, n_flank_list, c_flank_list)) + results = [] + for i, peptide in enumerate(peptide_list): + n_flank = n_flank_list[i] if n_flank_list is not None else "" + c_flank = c_flank_list[i] if c_flank_list is not None else "" + results.append(PeptideResult(preds=(Prediction( + kind=Kind.pMHC_presentation, + score=float(i), + peptide=peptide, + allele="HLA-A*02:01", + n_flank=n_flank, + c_flank=c_flank, + predictor_name="flank_echo", + ),))) + return results + + def predict_with_flanks(self, peptides, n_flanks, c_flanks): + return self.predict(peptides, n_flanks=n_flanks, c_flanks=c_flanks) + + +def test_predict_proteins_preserves_distinct_flanking_contexts(): + predictor = FlankEchoPredictor() + result = predictor.predict_proteins({"protein": "XABCYABCZ"}) + + peptides, n_flanks, c_flanks = predictor.calls[0] + assert peptides.count("ABC") == 2 + + abc_results = [pp for pp in result["protein"] if pp.peptide == "ABC"] + assert len(abc_results) == 2 + assert [(pp.offset, pp.presentation.n_flank, pp.presentation.c_flank) + for pp in abc_results] == [ + (1, "X", "YA"), + (5, "CY", "Z"), + ] + assert n_flanks[1] == "X" + assert c_flanks[1] == "YA" + assert n_flanks[5] == "CY" + assert c_flanks[5] == "Z" + + def test_predict_proteins_dataframe(): predictor = RandomBindingPredictor( alleles=["HLA-A*02:01"], From 67395f82b2fb7efc4f7438c07c108a0b6c815643 Mon Sep 17 00:00:00 2001 From: Alex Rubinsteyn Date: Wed, 6 May 2026 11:50:41 -0400 Subject: [PATCH 2/2] Relax MHCflurry parity tolerance --- tests/test_mhcflurry.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/test_mhcflurry.py b/tests/test_mhcflurry.py index 0e772d4..88b48c9 100644 --- a/tests/test_mhcflurry.py +++ b/tests/test_mhcflurry.py @@ -109,10 +109,14 @@ def test_mhcflurry_predict_proteins_matches_direct_flanked_prediction(): assert protein_result.presentation.c_flank == c_flank testing.assert_allclose( protein_result.presentation.score, - direct.presentation_score.iloc[0]) + direct.presentation_score.iloc[0], + rtol=1e-6, + atol=1e-6) testing.assert_allclose( protein_result.presentation.percentile_rank, - direct.presentation_percentile.iloc[0]) + direct.presentation_percentile.iloc[0], + rtol=1e-6, + atol=1e-6) def test_mhcflurry_presentation_affinity_matches_old_api():