Merge branch '372' into 431

healthcareai/common/model_eval.py

@@ -27,11 +27,14 @@ def compute_confusion_matrix(y_test, class_predictions):
     Returns:
         confusion_matrix: a confusion matrix
         class_names: names of the classes from the true label
+
+    Raises:
+        HealthcareAIError: Unequal number of predictions and true labels
     """
     _validate_predictions_and_labels_are_equal_length(class_predictions, y_test)
 
     # Get a list of classes in the test dataset
-    class_names = list(set(y_test))
+    class_names = sorted(set(y_test))
 
     # Compute confusion matrix
     confusion_matrix = skmetrics.confusion_matrix(

healthcareai/tests/test_model_eval.py

@@ -9,6 +9,41 @@
 import healthcareai.tests.helpers as test_helpers
 
 
+class TestConfusionMatrix(unittest.TestCase):
+    def test_raise_error_on_unequal_length(self):
+        truth = ['y', 'n', 'n', 'y', 'n']
+        preds = ['y', 'n', 'n', 'y']
+
+        self.assertRaises(
+            HealthcareAIError,
+            hcai_eval.compute_confusion_matrix,
+            truth,
+            preds)
+
+    def test_returns_correct_types(self):
+        truth = ['y', 'n', 'n', 'y']
+        preds = ['y', 'n', 'n', 'y']
+
+        cm, labels = hcai_eval.compute_confusion_matrix(truth, preds)
+
+        self.assertIsInstance(labels, list)
+        self.assertIsInstance(cm, np.ndarray)
+
+    def test_class_labels_sort(self):
+        truth_1 = ['y', 'y', 'n', 'y']
+        preds_1 = ['n', 'n', 'n', 'y']
+
+        truth_2 = ['n', 'n', 'n', 'y']
+        preds_2 = ['y', 'y', 'n', 'y']
+
+        _, obs_labels_1 = hcai_eval.compute_confusion_matrix(truth_1, preds_1)
+        _, obs_labels_2 = hcai_eval.compute_confusion_matrix(truth_2, preds_2)
+
+        print(obs_labels_1, '\n', obs_labels_2)
+
+        self.assertListEqual(obs_labels_1, obs_labels_2)
+
+
 class TestROC(unittest.TestCase):
     def test_roc(self):
         df = pd.DataFrame({'a': np.repeat(np.arange(.1, 1.1, .1), 10)})