Merge pull request #191 from bnaul/download_fix

Fix prediction downloads for probabilistic classifiers

Author: stefanv

cesium_app/handlers/prediction.py

@@ -141,8 +141,8 @@ def get(self, prediction_id=None, action=None):
                                    'label': data['labels'],
                                    'prediction': data['preds']},
                                   columns=['ts_name', 'label', 'prediction'])
-            if data.get('pred_probs'):
-                result['probability'] = np.max(data['pred_probs'], axis=1)
+            if len(data.get('pred_probs', [])) > 0:
+                result['probability'] = data['pred_probs'].max(axis=1).values
             self.set_header("Content-Type", 'text/csv; charset="utf-8"')
             self.set_header("Content-Disposition", "attachment; "
                             "filename=cesium_prediction_results.csv")

cesium_app/tests/frontend/test_predict.py

@@ -7,6 +7,7 @@
 from os.path import join as pjoin
 import numpy as np
 import numpy.testing as npt
+import pandas as pd
 from cesium_app.config import cfg
 import json
 import requests
@@ -185,6 +186,26 @@ def test_download_prediction_csv_class(driver):
             os.remove('/tmp/cesium_prediction_results.csv')
 
 
+def test_download_prediction_csv_class_prob(driver):
+    driver.get('/')
+    with create_test_project() as p, create_test_dataset(p) as ds,\
+         create_test_featureset(p) as fs,\
+         create_test_model(fs, model_type='RandomForestClassifier') as m,\
+         create_test_prediction(ds, m):
+        _click_download(p.id, driver)
+        assert os.path.exists('/tmp/cesium_prediction_results.csv')
+        try:
+            result = pd.read_csv('/tmp/cesium_prediction_results.csv')
+            npt.assert_array_equal(result.ts_name, np.arange(5))
+            npt.assert_array_equal(result.label, ['Mira', 'Classical_Cepheid',
+                                                  'Mira', 'Classical_Cepheid',
+                                                  'Mira'])
+            npt.assert_array_equal(result.label, result.prediction)
+            assert (result.probability >= 0.0).all()
+        finally:
+            os.remove('/tmp/cesium_prediction_results.csv')
+
+
 def test_download_prediction_csv_regr(driver):
     driver.get('/')
     with create_test_project() as p, create_test_dataset(p, label_type='regr') as ds,\