Files
vibrantlabsai--ragas/tests/unit/test_average_precision_algorithm.py
2026-07-13 13:35:10 +08:00

89 lines
3.0 KiB
Python

"""
Unit tests for Average Precision algorithm.
"""
from typing import List
import numpy as np
import pytest
def calculate_average_precision_original(verdict_list: List[int]) -> float:
"""Original implementation for comparison."""
if not verdict_list:
return 0.0
numerator = sum(
[
(sum(verdict_list[: i + 1]) / (i + 1)) * verdict_list[i]
for i in range(len(verdict_list))
]
)
denominator = sum(verdict_list) + 1e-10
return numerator / denominator
def calculate_average_precision_optimized(verdict_list: List[int]) -> float:
"""Optimized implementation matching the codebase."""
cumsum = 0
numerator = 0.0
for i, v in enumerate(verdict_list):
cumsum += v
if v:
numerator += cumsum / (i + 1)
denominator = cumsum + 1e-10
return numerator / denominator
class TestAveragePrecisionAlgorithm:
"""Test suite for Average Precision algorithm correctness."""
@pytest.mark.parametrize(
"verdict_list",
[
[], # empty
[1], # single positive
[0], # single negative
[1, 1, 1, 1, 1], # all ones
[0, 0, 0, 0, 0], # all zeros
[1, 0, 1], # alternating
[1, 1, 0, 1], # mixed
[0, 0, 1, 1, 1], # late positives
[1, 1, 0, 0, 1, 1, 0, 1], # realistic pattern
],
)
def test_optimized_matches_original(self, verdict_list):
"""Test that optimized algorithm produces identical results to original."""
original = calculate_average_precision_original(verdict_list)
optimized = calculate_average_precision_optimized(verdict_list)
assert np.isclose(original, optimized, rtol=1e-10, atol=1e-10)
def test_known_example_1_0_1(self):
"""Test [1,0,1]: score = (1 + 2/3) / 2 = 5/6."""
assert np.isclose(
calculate_average_precision_optimized([1, 0, 1]), 5 / 6, rtol=1e-10
)
def test_known_example_1_1_0_1(self):
"""Test [1,1,0,1]: score = (1 + 1 + 3/4) / 3 = 11/12."""
assert np.isclose(
calculate_average_precision_optimized([1, 1, 0, 1]), 11 / 12, rtol=1e-10
)
def test_early_positives_score_higher(self):
"""Earlier positives should score higher than later positives."""
early = calculate_average_precision_optimized([1, 1, 0, 0, 0])
late = calculate_average_precision_optimized([0, 0, 0, 1, 1])
assert early > late
@pytest.mark.parametrize("seed", [42, 123, 456])
def test_random_inputs(self, seed):
"""Test with random inputs for robustness."""
np.random.seed(seed)
for length in [10, 50, 100]:
verdict_list = np.random.choice([0, 1], size=length).tolist()
original = calculate_average_precision_original(verdict_list)
optimized = calculate_average_precision_optimized(verdict_list)
assert np.isclose(original, optimized, rtol=1e-10, atol=1e-10)