How to format xticklabels in a confusion matrix plotted with scikit-learn / matplotlib?

Question

I've plotted a confusion matrix with scikit-learn / matplotlib thanks to different code examples I found on the web, but I'm stuck at finding how to add space between the xticklabels and the main title. As you can see on the image below, the plot title and the xticklabels are overlapping (+ the ylabel 'True' is cut out).

Link to my confusion matrix image

Here is the function I use:

from sklearn.metrics import confusion_matrix
import matplotlib.pyplot as plt
PLOTS = '/plots/'  # Output folder

def plt_confusion_matrix(y_test, y_pred, normalize=False, title="Confusion matrix"):
    """
    Plots a nice confusion matrix.
    :param y_test: list of predicted labels
    :param y_pred: list of labels that should have been predicted.
    :param normalize: boolean. If False, the plots shows the number of sentences predicted.
    If True, shows the percentage of sentences predicted.
    :param title: string. Title of the plot.
    :return: Nothing but saves the plot as a PNG file and shows it.
    """
    labels = list(set(y_pred))
    cm = confusion_matrix(y_test, y_pred, labels)
    fig = plt.figure()
    ax = fig.add_subplot(111)

    cax = ax.matshow(cm, cmap=plt.cm.binary, interpolation='nearest')
    if normalize:
        cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]

    fig.suptitle(title, fontsize=14, wrap=True)
    fig.colorbar(cax)
    ax.set_xticklabels([''] + labels, rotation=45)
    ax.set_yticklabels([''] + labels)
    plt.xlabel('Predicted')
    plt.ylabel('True')
    plt.subplots_adjust(hspace=0.6)

    fmt = '.2f' if normalize else 'd'
    thresh = cm.max() / 1.5 if normalize else cm.max() / 2
    for i in range(cm.shape[0]):
        for j in range(cm.shape[1]):
            ax.text(j, i, format(cm[i, j], fmt),
                    ha="center", va="center",
                    color="white" if cm[i, j] > thresh else "black")

    plt.savefig(PLOTS + title)
    plt.show()

I had to rotate the xticklabels because they are too long and otherwise overlapping each other horizontally, and I had to wrap the title because it is also too long and could not otherwise be displayed entirely in the image.

I've seen in another post that xticklabels can also be placed underneath the figure (like in this stackoverflow post), so maybe it could be a solution, however I haven't understood how to make it.

How do I solve the problem?

• either to add some space between the title and the xticklabels (making them appear entirely btw);
• or to make the ylabel 'True' visible
• or to move the xticklabels under the figure.

Edit : I tried both of geekzeus solutions, without success...

• Result with geekzeus' 1st solution : See confusion matrix
• Result with geekzeus' 2nd solution : See confusion matrix

Answer 1

Do it like this

ax.set_xlabel('Predicted labels') 
ax.set_ylabel('True labels')       
ax.set_title('Confusion Matrix')   
#xaxisticks
ax.xaxis.set_ticklabels(['A', 'B']) 
#yaxis ticks
ax.yaxis.set_ticklabels(['B', 'A']) 

OR use seaborn with matplotlib,you can also directly provide list variable to ticks

import seaborn as sns
import matplotlib.pyplot as plt     

cm = confusion_matrix(true_classes, predicted_classes)

ax= plt.subplot()
sns.heatmap(cm, annot=True, ax = ax); #annot=True to annotate cells

# labels, title and ticks
ax.set_xlabel('Predicted labels')
ax.set_ylabel('True labels') 
ax.set_title('Confusion Matrix')
ax.xaxis.set_ticklabels(['A', 'B'])   
ax.yaxis.set_ticklabels(['B', 'A'])

Answer 2

You can specify the location of the title using parameters x and y. If you tweak the values of y, the desired plot can be generated.

fig.suptitle(title, fontsize=14, wrap=True, y=1.2)