varunagrawal · December 10, 2018 23:24 · Dec 10, 2018
diff --git a/average_precision.py b/average_precision.py
@@ -0,0 +1,38 @@
+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