add hide_counts parameter to plot_confusion_matrix() (#90)

* add hide_counts parameter to plot_confusion_matrix()

* add test for hide_counts argument in plot_confusion_matrix()

Author: echan5

scikitplot/metrics.py

@@ -32,7 +32,7 @@
 
 def plot_confusion_matrix(y_true, y_pred, labels=None, true_labels=None,
                           pred_labels=None, title=None, normalize=False,
-                          hide_zeros=False, x_tick_rotation=0, ax=None,
+                          hide_zeros=False, hide_counts=False, x_tick_rotation=0, ax=None,
                           figsize=None, cmap='Blues', title_fontsize="large",
                           text_fontsize="medium"):
     """Generates confusion matrix plot from predictions and true labels
@@ -65,6 +65,9 @@ def plot_confusion_matrix(y_true, y_pred, labels=None, true_labels=None,
         hide_zeros (bool, optional): If True, does not plot cells containing a
             value of zero. Defaults to False.
 
+        hide_counts (bool, optional): If True, doe not overlay counts.
+            Defaults to False.
+
         x_tick_rotation (int, optional): Rotates x-axis tick labels by the
             specified angle. This is useful in cases where there are numerous
             categories and the labels overlap each other.
@@ -160,13 +163,15 @@ def plot_confusion_matrix(y_true, y_pred, labels=None, true_labels=None,
     ax.set_yticklabels(true_classes, fontsize=text_fontsize)
 
     thresh = cm.max() / 2.
-    for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):
-        if not (hide_zeros and cm[i, j] == 0):
-            ax.text(j, i, cm[i, j],
-                    horizontalalignment="center",
-                    verticalalignment="center",
-                    fontsize=text_fontsize,
-                    color="white" if cm[i, j] > thresh else "black")
+
+    if not hide_counts:
+        for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):
+            if not (hide_zeros and cm[i, j] == 0):
+                ax.text(j, i, cm[i, j],
+                        horizontalalignment="center",
+                        verticalalignment="center",
+                        fontsize=text_fontsize,
+                        color="white" if cm[i, j] > thresh else "black")
 
     ax.set_ylabel('True label', fontsize=text_fontsize)
     ax.set_xlabel('Predicted label', fontsize=text_fontsize)

scikitplot/tests/test_metrics.py

@@ -59,6 +59,13 @@ def test_labels(self):
         preds = clf.predict(self.X)
         plot_confusion_matrix(self.y, preds, labels=[0, 1, 2])
 
+    def test_hide_counts(self):
+        np.random.seed(0)
+        clf = LogisticRegression()
+        clf.fit(self.X, self.y)
+        preds = clf.predict(self.X)
+        plot_confusion_matrix(self.y, preds, hide_counts=True)
+
     def test_true_pred_labels(self):
         np.random.seed(0)
         clf = LogisticRegression()