Search code examples
pythonpandasseaborndiagonal

How to sort a confusion matrix by the diagonal value


I have this confusion matrix:

import pandas as pd
import seaborn as sn
import matplotlib.pyplot as plt

data = {'y_Actual':    [3, 3, 1, 1, 0, 1, 2, 3, 1, 1, 1, 0, 2, 4, 3],
        'y_Predicted': [1, 2, 2, 1, 0, 1, 3, 0, 1, 0, 0, 0, 3, 4, 2]
        }

df = pd.DataFrame(data, columns=['y_Actual','y_Predicted'])
confusion_matrix = pd.crosstab(df['y_Actual'], df['y_Predicted'],
                               rownames=['Actual'], colnames=['Predicted'])

sn.heatmap(confusion_matrix, annot=True)
plt.show()

enter image description here

Is it possible to sort this confusion matrix by diagonal value in descending order ?


Solution

  • You can use numpy.argsort on the diagonal (opposite values for descending order) and reindex:

    import numpy as np
    order = np.argsort(-confusion_matrix.to_numpy().diagonal())
    # array([1, 0, 4, 2, 3])
    
    sn.heatmap(confusion_matrix.iloc[order, order], annot=True)
    

    output:

    enter image description here