diff --git a/mhcflurry/class1_affinity_prediction/class1_affinity_predictor.py b/mhcflurry/class1_affinity_prediction/class1_affinity_predictor.py index 5d6df8c6d388fb488f1961db7dbc1aeed8b17ce0..b424aaf5bae2b75a0d86c0a1be6b92427faa5cdf 100644 --- a/mhcflurry/class1_affinity_prediction/class1_affinity_predictor.py +++ b/mhcflurry/class1_affinity_prediction/class1_affinity_predictor.py @@ -753,7 +753,11 @@ class Class1AffinityPredictor(object): mask = ( (df.normalized_allele == allele) & df.supported_peptide_length).values - if mask.sum() > 0: + if mask.all(): + # Common case optimization + for (i, model) in enumerate(models): + df["model_single_%d" % i] = model.predict(peptides) + elif mask.sum() > 0: allele_peptides = EncodableSequences.create( df.ix[mask].peptide.values) for (i, model) in enumerate(models):