Update ap.py

foryouforever · Jun 29, 2017 · d2508a8 · d2508a8
1 parent b0ca6fe
commit d2508a8
Showing 1 changed file with 11 additions and 11 deletions.
diff --git a/pyltr/metrics/ap.py b/pyltr/metrics/ap.py
@@ -26,7 +26,7 @@ def evaluate(self, qid, targets):
                 total_prec += num_rel / (i + 1.0)
         return (total_prec / num_rel) if num_rel > 0 else 0.0
 
-    def calc_swap_deltas(self, qid, targets):
+def calc_swap_deltas(self, qid, targets):
         n_targets = len(targets)
         deltas = np.zeros((n_targets, n_targets))
         total_num_rel = 0
@@ -42,18 +42,16 @@ def calc_swap_deltas(self, qid, targets):
             if targets[i] >= self.cutoff:
                 num_rel_i += 1
                 num_rel_j = num_rel_i
-                sub = 1 / (i + 1.0)
+                sub = num_rel_i / (i + 1.0)
 
                 for j in range(i + 1, n_targets):
                     if targets[j] >= self.cutoff:
                         if j < self.k:
                             num_rel_j += 1
                             sub += 1 / (j + 1.0)
                     else:
-                        #add = (num_rel_j / (j + 1.0)) if j < self.k else 0.0
-                        new_total_metric = total_metric - sub
-                        if j < self.k:
-                            sub += 1 / (j + 1.0)
+                        add = (num_rel_j / (j + 1.0)) if j < self.k else 0.0
+                        new_total_metric = total_metric + add - sub
                         new_num_rel = (total_num_rel
                                        if j < self.k
                                        else (total_num_rel - 1))
@@ -64,10 +62,14 @@ def calc_swap_deltas(self, qid, targets):
 
             else:
                 num_rel_j = num_rel_i
-                add = 1 / (i + 1.0)
+                add = (num_rel_i + 1) / (i + 1.0)
+
                 for j in range(i + 1, n_targets):
                     if targets[j] >= self.cutoff:
-                        new_total_metric = total_metric + add
+                        sub = (((num_rel_j + 1) / (j + 1.0))
+                               if j < self.k
+                               else 0.0)
+                        new_total_metric = total_metric + add - sub
                         new_num_rel = (total_num_rel
                                        if j < self.k
                                        else (total_num_rel + 1))
@@ -79,9 +81,7 @@ def calc_swap_deltas(self, qid, targets):
                         if j < self.k:
                             num_rel_j += 1
                             add += 1 / (j + 1.0)
-                    else:
-                        if j < self.k:
-                            add += 1 / (j + 1.0)
+
         return deltas
 
     def max_k(self):