fix

mhcflurry/antigen_presentation/presentation_component_models/presentation_component_model.py

@@ -233,7 +233,8 @@ class PresentationComponentModel(object):
             return_value = result_df[columns]
         if self.cached_predictions is not None:
             self.cached_predictions[cache_key] = return_value
-        return return_value
+        return dict(
+            (col, return_value[col].values) for col in self.column_names())
 
     def clone(self):
         """

mhcflurry/antigen_presentation/presentation_model.py

@@ -188,7 +188,7 @@ class PresentationModel(object):
                         model_input_training_hits_df)
                     self.trained_component_models[-1].append((sub_model,))
                     predictions = sub_model.predict(hits_and_decoys_df)
-                    for (col, values) in predictions.iteritems():
+                    for (col, values) in predictions.items():
                         hits_and_decoys_df[col] = values
                 final_predictor = self.fit_final_predictor(hits_and_decoys_df)
                 self.presentation_models_predictors.append(final_predictor)