Resolve "Add survival tests for all plotting methods"

Closes #22 (closed)

nnfwtbn/tests/test_plot.py

@@ -7,7 +7,7 @@ import matplotlib.pyplot as plt
 from nnfwtbn.variable import Variable
 from nnfwtbn.cut import Cut
 from nnfwtbn.process import Process
-from nnfwtbn.plot import roc
+from nnfwtbn.plot import roc, confusion_matrix
 
 
 class SurvivalTestCase(unittest.TestCase):
@@ -68,3 +68,27 @@ class SurvivalTestCase(unittest.TestCase):
         except Exception:
             self.fail("Calling roc() failed.")
 
+
+    def test_confusion_matrix(self):
+        """
+        Check that calling confusion_matrix() does not raise an exception.
+        """
+        df = pd.DataFrame({
+                "sick": [0, 1, 1, 0, 0],
+                "positive": [0, 0, 1, 1, 0],
+                "weight": [1, 1, 1, 1.7, 1],
+            })
+
+        positive = Cut(lambda d: d.positive == 1, label="+")
+        negative = Cut(lambda d: d.positive == 0, label="-")
+        sick = Cut(lambda d: d.sick == 1, label="sick")
+        healthy = Cut(lambda d: d.sick == 0, label="healthy")
+
+        try:
+            data = confusion_matrix(df, 
+                             [positive, negative], [sick, healthy],
+                             "Test result", "Truth",
+                             weight=Variable("weight", "weight"),
+                             annot=True)
+        except Exception:
+            self.fail("Calling confusion_matrix() failed.")