Search code examples
pythonmachine-learningscikit-learnconfusion-matrix

How to plot 2x2 confusion matrix with predictions in rows an real values in columns?


I know that we can plot a confusion matrix with sklearn using the following sample code.

from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay
import matplotlib.pyplot as plt

y_true = [1, 0, 1, 1, 0, 1]
y_pred = [0, 0, 1, 1, 0, 1]

print(f'y_true: {y_true}')
print(f'y_pred: {y_pred}\n')

cm = confusion_matrix(y_true, y_pred, labels=[0, 1])
print(cm)
disp = ConfusionMatrixDisplay(confusion_matrix=cm)
disp.plot()
plt.show()

enter image description here

What we have:

TN | FP
FN | TP

But I want the prediction label placed in a row or y-axis and the true or real value label in a column or x-axis. How can I plot this using Python?

What I want:

TP | FP
FN | TN

Solution

  • (1) Here is one way of reversing TP/TN.

    Code

    """
    Reverse True and Prediction labels
    
    References:
        https://github.com/scikit-learn/scikit-learn/blob/0d378913b/sklearn/metrics/_plot/confusion_matrix.py
        https://scikit-learn.org/stable/modules/generated/sklearn.metrics.ConfusionMatrixDisplay.html
    """
    
    from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay
    import matplotlib.pyplot as plt
    
    y_true = [1, 0, 1, 1, 0, 1]
    y_pred = [0, 0, 1, 1, 0, 1]
    
    print(f'y_true: {y_true}')
    print(f'y_pred: {y_pred}\n')
    
    # Normal
    print('Normal')
    cm = confusion_matrix(y_true, y_pred, labels=[0, 1])
    print(cm)
    disp = ConfusionMatrixDisplay(confusion_matrix=cm)
    disp.plot()
    
    plt.savefig('normal.png')
    plt.show()
    
    # Reverse TP and TN
    print('Reverse TP and TN')
    cm = confusion_matrix(y_pred, y_true, labels=[1, 0])  # reverse true/pred and label values
    print(cm)
    disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=[1, 0])  # reverse display labels
    dp = disp.plot()
    dp.ax_.set(ylabel="My Prediction Label")  # modify ylabel of ax_ attribute of plot
    dp.ax_.set(xlabel="My True Label")        # modify xlabel of ax_ attribute of plot
    
    plt.savefig('reverse.png')
    plt.show()
    

    Output

    y_true: [1, 0, 1, 1, 0, 1]
    y_pred: [0, 0, 1, 1, 0, 1]
    
    Normal
    [[2 0]
     [1 3]]
    

    enter image description here

    Reverse TP and TN
    [[3 0]
     [1 2]]
    

    enter image description here

    (2) Another way is by swapping values and plot it with sns/matplotlib.

    Code

    import seaborn as sns
    from sklearn.metrics import confusion_matrix
    import matplotlib.pyplot as plt
    
    
    y_true = [1, 0, 1, 1, 0, 1]
    y_pred = [0, 0, 1, 1, 0, 1]
    
    cm = confusion_matrix(y_true, y_pred)
    print(cm)
    cm_11 = cm[1][1]     # backup value in cm[1][1]
    cm[1][1] = cm[0][0]  # swap
    cm[0][0] = cm_11     # swap
    print(cm)
    
    ax = sns.heatmap(cm, annot=True)
    
    plt.yticks([1.5, 0.5], ['0', '1'], ha='right')
    plt.xticks([1.5, 0.5], ['0', '1'], ha='right')
    
    ax.set(xlabel='True Label', ylabel='Prediction Label')
    plt.savefig('reverse_tp_tn.png')
    plt.show()
    

    Output

    [[2 0]
     [1 3]]
    [[3 0]
     [1 2]]
    

    enter image description here