average_precision.py

import numpy as np


def average_precision(recall, precision):
    mrec = np.hstack((0, recall, 1))
    mpre = np.hstack((0, precision, 0))

    for i in range(mpre.size-2, -1, -1):
        mpre[i] = max(mpre[i], mpre[i+1])

    i = np.ravel(np.where(mrec[1:] != mrec[0:-1])) + 1

    ap = np.sum((mrec[i]-mrec[i-1]) * mpre[i])
    return ap


def precision_recall(truth, scores, pos_label=1, neg_label=0):
    desc_score_idx = np.argsort(-scores, kind='stable')
    scores = scores[desc_score_idx]
    truth = truth[desc_score_idx]

    distinct_value_indices = np.where(np.diff(scores))[0]
    threshold_idxs = np.r_[distinct_value_indices, truth.size - 1]

    tp = (truth == pos_label).astype(np.float)
    fp = (truth == neg_label).astype(np.float)

    tps = np.cumsum(tp)[threshold_idxs]
    fps = np.cumsum(fp)[threshold_idxs]

    precision = tps / (tps + fps)

    if tps[-1] == 0:
        recall = np.ones(tps.size)
    else:
        recall = tps / tps[-1]

    return precision, recall