diff --git a/great_ai/utilities/evaluate_ranking/evaluate_ranking.py b/great_ai/utilities/evaluate_ranking/evaluate_ranking.py index a613c93..bb6ce34 100644 --- a/great_ai/utilities/evaluate_ranking/evaluate_ranking.py +++ b/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): diff --git a/tests/utilities/test_evaluate_ranking.py b/tests/utilities/test_evaluate_ranking.py index 8111b6a..23cf14d 100644 --- a/tests/utilities/test_evaluate_ranking.py +++ b/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: