diff --git a/mhcflurry/class1_affinity_predictor.py b/mhcflurry/class1_affinity_predictor.py index 4a4144ea6ba7b9dbd1491c1b24346390e8cb080f..4773df956e4cdbfbd3dbce8a4eb26446f39359d0 100644 --- a/mhcflurry/class1_affinity_predictor.py +++ b/mhcflurry/class1_affinity_predictor.py @@ -1046,4 +1046,31 @@ class Class1AffinityPredictor(object): return encoded_peptides + def filter_networks(self, predicate): + """ + Return a new Class1AffinityPredictor containing a subset of this + predictor's neural networks. + + Parameters + ---------- + predicate : Class1NeuralNetwork -> boolean + Function specifying which neural networks to include + Returns + ------- + Class1AffinityPredictor + """ + allele_to_allele_specific_models = {} + for (allele, models) in self.allele_to_allele_specific_models.items(): + allele_to_allele_specific_models[allele] = [ + m for m in models if predicate(m) + ] + class1_pan_allele_models = [ + m for m in self.class1_pan_allele_models if predicate(m) + ] + + return Class1AffinityPredictor( + allele_to_allele_specific_models=allele_to_allele_specific_models, + class1_pan_allele_models=class1_pan_allele_models, + allele_to_fixed_length_sequence=self.allele_to_fixed_length_sequence, + ) \ No newline at end of file