diff --git a/mhcflurry/antigen_presentation/presentation_model.py b/mhcflurry/antigen_presentation/presentation_model.py index 73c1f68832fe464f24491104850cd9babae51fa0..33e3e7099d9a9bc49fdf4e9eee09054932cc96dd 100644 --- a/mhcflurry/antigen_presentation/presentation_model.py +++ b/mhcflurry/antigen_presentation/presentation_model.py @@ -422,6 +422,8 @@ class PresentationModel(object): assert 'experiment_name' in peptides_df.columns assert 'hit' in peptides_df.columns + peptides_df = peptides_df.copy() + peptides_df["prediction"] = self.predict(peptides_df) top_n = float(peptides_df.hit.sum()) @@ -432,8 +434,13 @@ class PresentationModel(object): } else: ranks = peptides_df.prediction.rank(ascending=False) + hit_indices = ranks[peptides_df.hit > 0].values + hit_lengths = peptides_df.peptide[ + peptides_df.hit > 0 + ].str.len().values result = { - 'hit_indices': numpy.sort(ranks[peptides_df.hit > 0].values), + 'hit_indices': hit_indices, + 'hit_lengths': hit_lengths, 'total_peptides': len(peptides_df), } result['score'] = (