ENH cosmetic reorg of the confusion matrix example

Author: ogrisel

examples/model_selection/plot_confusion_matrix.py

@@ -27,13 +27,12 @@
 print(__doc__)
 
 import numpy as np
+import matplotlib.pyplot as plt
 
 from sklearn import svm, datasets
 from sklearn.cross_validation import train_test_split
 from sklearn.metrics import confusion_matrix
 
-import matplotlib.pyplot as plt
-
 # import some data to play with
 iris = datasets.load_iris()
 X = iris.data
@@ -47,40 +46,33 @@
 classifier = svm.SVC(kernel='linear', C=0.01)
 y_pred = classifier.fit(X_train, y_train).predict(X_test)
 
+
+def plot_confusion_matrix(cm, title='Confusion matrix', cmap=plt.cm.Blues):
+    plt.imshow(cm, interpolation='nearest', cmap=cmap)
+    plt.title(title)
+    plt.colorbar()
+    tick_marks = np.arange(len(iris.target_names))
+    plt.xticks(tick_marks, iris.target_names, rotation=45)
+    plt.yticks(tick_marks, iris.target_names)
+    plt.tight_layout()
+    plt.ylabel('True label')
+    plt.xlabel('Predicted label')
+
+
 # Compute confusion matrix
 cm = confusion_matrix(y_test, y_pred)
+np.set_printoptions(precision=2)
 print('Confusion matrix, without normalization')
 print(cm)
-
-# Show confusion matrix in a separate window
-plt.imshow(cm, interpolation='nearest', cmap=plt.cm.binary)
-plt.title('Confusion matrix')
-plt.set_cmap('Blues')
-plt.colorbar()
-tick_marks = np.arange(len(iris.target_names))
-plt.xticks(tick_marks, iris.target_names, rotation=60)
-plt.yticks(tick_marks, iris.target_names)
-plt.ylabel('True label')
-plt.xlabel('Predicted label')
-# Convenience function to adjust plot parameters for a clear layout.
-plt.tight_layout()
+plt.figure()
+plot_confusion_matrix(cm)
 
 # Normalize the confusion matrix by row (i.e by the number of samples
 # in each class)
 cm_normalized = cm.astype('float') / cm.sum(axis=1)
-
 print('Normalized confusion matrix')
 print(cm_normalized)
-
-# Show normalized confusion matrix in a separate window
 plt.figure()
-plt.imshow(cm_normalized, interpolation='nearest', cmap=plt.cm.binary)
-plt.title('Normalized confusion matrix')
-plt.set_cmap('Blues')
-plt.colorbar()
-plt.xticks(tick_marks, iris.target_names, rotation=60)
-plt.yticks(tick_marks, iris.target_names)
-plt.ylabel('True label')
-plt.xlabel('Predicted label')
-plt.tight_layout()
+plot_confusion_matrix(cm_normalized, title='Normalized confusion matrix')
+
 plt.show()