Calculate ranking recall from relevant items

great_ai/utilities/evaluate_ranking/evaluate_ranking.py

@@ -4,7 +4,7 @@ from typing import Dict, List, Optional, TypeVar
 import matplotlib
 import matplotlib.pyplot as plt
 import numpy as np
-from sklearn.metrics import average_precision_score
+from sklearn.metrics import average_precision_score, precision_recall_curve
 
 from ..unique import unique
 from .draw_f1_iso_lines import draw_f1_iso_lines
@@ -68,16 +68,13 @@ def evaluate_ranking(
                 for v in expected
             ]
 
-            sorted_expected_actual = sorted(
-                zip(binarized_expected, actual_scores), key=lambda v: v[1], reverse=True
+            precision_values, recall_values, _ = precision_recall_curve(
+                binarized_expected, actual_scores
             )
-            precision = []
-            recall = []
-            correct = 0
-            for all, (e, score) in enumerate(sorted_expected_actual, start=1):
-                correct += int(e)
-                precision.append(correct / all)
-                recall.append(all / len(sorted_expected_actual))
+            # Score thresholds keep tied scores together. Exclude the synthetic
+            # empty-prediction endpoint and order by increasing recall.
+            precision = precision_values[:-1][::-1].tolist()
+            recall = recall_values[:-1][::-1].tolist()
 
             if not disable_interpolation:
                 for j in range(len(precision) - 2, -1, -1):

tests/utilities/test_evaluate_ranking.py

@@ -16,7 +16,7 @@ def test_default() -> None:
         reverse_order=True,
     )
 
-    assert results == {"d": 5 / 6, "c": 2 / 3, "b": 2 / 5}
+    assert results == {"d": 1.0, "c": 1.0, "b": 1 / 2}
 
     results = evaluate_ranking(
         ["a", "a", "b", "b", "c", "d", "d", "d"],
@@ -26,7 +26,23 @@ def test_default() -> None:
         disable_interpolation=True,
     )
 
-    assert results == {"a": 0.6, "b": 0.4, "c": 0.2}
+    assert results == {"a": 2 / 3, "b": 2 / 5, "c": 2 / 7}
+
+
+@pytest.mark.parametrize("disable_interpolation", [False, True])
+def test_perfect_ranking_at_full_recall(disable_interpolation: bool) -> None:
+    assert evaluate_ranking(
+        [1, 1, 0, 0],
+        [4, 3, 2, 1],
+        target_recall=1,
+        disable_interpolation=disable_interpolation,
+        plot=False,
+    ) == {0: 1.0}
+
+
+def test_tied_scores_cannot_be_separated_by_a_threshold() -> None:
+    for labels in ([1, 0], [0, 1]):
+        assert evaluate_ranking(labels, [1, 1], target_recall=1, plot=False) == {0: 0.5}
 
 
 def test_mismatching_lengths() -> None: