How to calculate TP/TN/FP/FN from binary value pairs in Python?

Question

I want to write a Python function that on a binary value input pair (truth, prediction) gives back True Positive/True Negative/False Positive/False Negative values according their input. So far I have reached the required output with this:

def func(truth, prediction):
    
    if prediction == 1:
        
        if truth == 1:
            return "TP"
        else:
            return "FP"

    elif truth == 1:
        return "FN"

    else:
        return "TN"

However, this seems a but clunky solution, is there a shorter, more elegant way?

(The input pair is supposed to be a binary integer 0/1)

Answer 1

The comment from Johnny Mopp is pretty cool (though I think the order should be ['TN', 'FN', 'FP', 'TP'] , but if I came across it in code I'd have to think twice. (I can count on one hand the number of times I've seen bit-shift operations in production code.)

There's a new way to handle things like this in Python 3.10: structural pattern matching . Now, this is the first time I've tried to use this new feature, but here's how it might look:

def get_result(true, pred):
    """Decide if TP, FP, TN or FN"""
    match [true, pred]:
        case [1, 1]: return 'TP'
        case [1, 0]: return 'FN'
        case [0, 0]: return 'TN'
        case [0, 1]: return 'FP'

Seems to work:

>>> y = [1, 0, 1, 0]  # TP, FP, FN, TN
>>> ŷ = [1, 1, 0, 0]

>>> for yi, ŷi in zip(y, ŷ):
>>>     print(get_result(yi, ŷi))
TP
FP
FN
TN

If you need to use Python 3.9 or below, then you could compactify your approach with something like this:

def get_result(true, pred):
    """Decide if TP, FP, TN or FN"""
    if pred:
        return 'TP' if true else 'FP'
    return 'FN' if true else 'TN'