a pandas wrapper for the sklearn confusion_matrix method

confusion_dataframe.py

"""A simple method wrapping around `confusion_matrix`."""


def confusion_dataframe(*args, **kwargs):
    """
    Generate an sklearn confusion matrix as a DataFrame.

    Args:
        *args: positional args for `confusion_matrix`
        **kwargs: keyword args for `confusion_matrix`

    Returns: a confusion matrix formatted into a pandas DataFrame

    """
    import pandas as pd
    from sklearn.metrics import confusion_matrix
    # generate the confusion matrix with the given arguments
    mat = confusion_matrix(*args, **kwargs)
    # convert the confusion matrix to a pandas dataframe
    df = pd.DataFrame(mat, columns=['N', 'P'], index=["N", "P"])
    # label the x,y indexes
    df.columns.name = 'Prediction'
    df.index.name = 'Actual'

    return df


__all__ = ['confusion_dataframe']