Recovered source transcript
1package eval
2
3import (
4 "math"
5 "testing"
6)
7
8// closeEnough compares two floats to a tight tolerance, since these metrics involve log2 divisions whose exact
9// decimal expansion is irrational. The tolerance is far tighter than any band a regression could move.
10func closeEnough(a, b float64) bool { return math.Abs(a-b) < 1e-9 }
11
12func rel(ids ...string) map[string]bool {
13 m := make(map[string]bool, len(ids))
14 for _, id := range ids {
15 m[id] = true
16 }
17 return m
18}
19
20// ---------- RecallAtK ----------
21
22func TestRecallAtK(t *testing.T) {
23 tests := []struct {
24 name string
25 ranked []string
26 relevant map[string]bool
27 k int
28 want float64
29 }{
30 // Perfect ranking: both relevant ids are in the top two, so recall at any k that reaches them is 1.0.
31 {"perfect both in top k", []string{"a", "b", "c"}, rel("a", "b"), 2, 1.0},
32 // Reversed ranking: the relevant ids sit at the bottom; a small k misses them entirely.
33 {"reversed misses at small k", []string{"c", "d", "a", "b"}, rel("a", "b"), 2, 0.0},
34 // Reversed but k large enough to reach both relevant ids.
35 {"reversed found at large k", []string{"c", "d", "a", "b"}, rel("a", "b"), 4, 1.0},
36 // Partial: one of two relevant ids inside the cutoff.
37 {"partial one of two", []string{"a", "x", "y"}, rel("a", "b"), 3, 0.5},
38 // No relevant id is present anywhere in the ranked list.
39 {"none relevant found", []string{"x", "y", "z"}, rel("a", "b"), 3, 0.0},
40 // k larger than the list: capped at the list length, both relevant ids still found.
41 {"k larger than list", []string{"a", "b"}, rel("a", "b"), 50, 1.0},
42 // Empty ranked list: nothing can be recalled.
43 {"empty ranked list", []string{}, rel("a"), 5, 0.0},
44 // Empty relevant set is defined as 1.0 (nothing was missed, vacuously).
45 {"empty relevant set", []string{"a", "b"}, rel(), 5, 1.0},
46 // k of zero recalls nothing.
47 {"k zero", []string{"a", "b"}, rel("a"), 0, 0.0},
48 // Negative k recalls nothing.
49 {"k negative", []string{"a", "b"}, rel("a"), -3, 0.0},
50 // Ties / duplicates: a repeated relevant id is counted once, so recall cannot exceed 1.0.
51 {"duplicate relevant counted once", []string{"a", "a", "a"}, rel("a", "b"), 3, 0.5},
52 // Single relevant id at the cutoff edge.
53 {"relevant at k boundary", []string{"x", "a"}, rel("a"), 2, 1.0},
54 // Single relevant id just past the cutoff edge.
55 {"relevant just past k", []string{"x", "a"}, rel("a"), 1, 0.0},
56 }
57 for _, tt := range tests {
58 t.Run(tt.name, func(t *testing.T) {
59 got := RecallAtK(tt.ranked, tt.relevant, tt.k)
60 if !closeEnough(got, tt.want) {
61 t.Fatalf("RecallAtK(%v, k=%d) = %v, want %v", tt.ranked, tt.k, got, tt.want)
62 }
63 })
64 }
65}
66
67// ---------- MRR ----------
68
69func TestMRR(t *testing.T) {
70 third := 1.0 / 3.0
71 tests := []struct {
72 name string
73 ranked []string
74 relevant map[string]bool
75 want float64
76 }{
77 // First relevant id at rank one scores the maximum.
78 {"first relevant at rank 1", []string{"a", "b", "c"}, rel("a"), 1.0},
79 // First relevant id at rank two scores one half.
80 {"first relevant at rank 2", []string{"x", "a", "c"}, rel("a"), 0.5},
81 // First relevant id at rank three scores one third.
82 {"first relevant at rank 3", []string{"x", "y", "a"}, rel("a"), third},
83 // Multiple relevant ids: only the FIRST one counts, the earliest rank wins.
84 {"earliest of several", []string{"x", "b", "a"}, rel("a", "b"), 0.5},
85 // No relevant id present anywhere.
86 {"none relevant", []string{"x", "y", "z"}, rel("a"), 0.0},
87 // Reversed ranking pushes the only relevant id to the bottom.
88 {"reversed to bottom", []string{"d", "c", "b", "a"}, rel("a"), 0.25},
89 // Empty ranked list.
90 {"empty ranked", []string{}, rel("a"), 0.0},
91 // Empty relevant set: there is no first relevant id, so MRR is 0.0.
92 {"empty relevant", []string{"a", "b"}, rel(), 0.0},
93 // Perfect: relevant id leads the list.
94 {"perfect lead", []string{"a"}, rel("a"), 1.0},
95 }
96 for _, tt := range tests {
97 t.Run(tt.name, func(t *testing.T) {
98 got := MRR(tt.ranked, tt.relevant)
99 if !closeEnough(got, tt.want) {
100 t.Fatalf("MRR(%v) = %v, want %v", tt.ranked, got, tt.want)
101 }
102 })
103 }
104}
105
106// ---------- NDCGAtK ----------
107
108func TestNDCGAtK(t *testing.T) {
109 // Hand computed reference values. The discount for the id at one based rank r is 1/log2(r+1).
110 // rank 1 discount = 1/log2(2) = 1.0
111 // rank 2 discount = 1/log2(3) = 0.6309297535714574
112 // rank 3 discount = 1/log2(4) = 0.5
113 d1, d2, d3 := 1.0, 1.0/math.Log2(3), 0.5
114
115 tests := []struct {
116 name string
117 ranked []string
118 relevant map[string]bool
119 k int
120 want float64
121 }{
122 // Perfect ranking: relevant ids fill the top, nDCG is 1.0 regardless of their internal order.
123 {"perfect packed top", []string{"a", "b", "c"}, rel("a", "b"), 3, 1.0},
124 {"perfect internal order irrelevant", []string{"b", "a", "c"}, rel("a", "b"), 3, 1.0},
125 // One relevant id at rank one with one irrelevant: ideal is also one relevant at rank one, so 1.0.
126 {"single relevant at top", []string{"a", "x"}, rel("a"), 2, 1.0},
127 // One relevant id at rank two: dcg = d2, idcg = d1, so nDCG = d2/d1 = d2.
128 {"single relevant rank 2", []string{"x", "a"}, rel("a"), 2, d2 / d1},
129 // One relevant id at rank three (k high enough): dcg = d3, idcg = d1.
130 {"single relevant rank 3", []string{"x", "y", "a"}, rel("a"), 3, d3 / d1},
131 // Two relevant ids at ranks one and three: dcg = d1 + d3, idcg = d1 + d2 (two relevant ideal).
132 {"two relevant ranks 1 and 3", []string{"a", "x", "b"}, rel("a", "b"), 3, (d1 + d3) / (d1 + d2)},
133 // Reversed: relevant ids at the bottom, both inside k.
134 {"reversed both in k", []string{"x", "y", "b", "a"}, rel("a", "b"), 4,
135 (1.0/math.Log2(4) + 1.0/math.Log2(5)) / (d1 + d2)},
136 // No relevant id in the top k.
137 {"none in top k", []string{"x", "y", "z"}, rel("a", "b"), 3, 0.0},
138 // k larger than the list: only what is present counts; both relevant ids are at the top, so 1.0.
139 {"k larger than list", []string{"a", "b"}, rel("a", "b"), 50, 1.0},
140 // k smaller than the relevant count: ideal is bounded by k, one relevant fits, so a single relevant
141 // id at rank one is perfect for that cutoff.
142 {"k below relevant count perfect", []string{"a", "x"}, rel("a", "b"), 1, 1.0},
143 // Empty ranked list.
144 {"empty ranked", []string{}, rel("a"), 5, 0.0},
145 // Empty relevant set is defined as a perfect, vacuous 1.0.
146 {"empty relevant", []string{"a", "b"}, rel(), 5, 1.0},
147 // k of zero yields 0.0 when there is something relevant to find.
148 {"k zero", []string{"a"}, rel("a"), 0, 0.0},
149 // Negative k yields 0.0.
150 {"k negative", []string{"a"}, rel("a"), -2, 0.0},
151 // Ties / duplicates: a repeated relevant id is credited only on its first appearance, so a second
152 // copy at rank two adds nothing; dcg = d1, idcg = d1 (one distinct relevant), nDCG = 1.0.
153 {"duplicate credited once", []string{"a", "a", "x"}, rel("a"), 3, 1.0},
154 }
155 for _, tt := range tests {
156 t.Run(tt.name, func(t *testing.T) {
157 got := NDCGAtK(tt.ranked, tt.relevant, tt.k)
158 if !closeEnough(got, tt.want) {
159 t.Fatalf("NDCGAtK(%v, k=%d) = %v, want %v", tt.ranked, tt.k, got, tt.want)
160 }
161 })
162 }
163}
164
165// TestNDCGBoundedZeroOne proves nDCG never escapes the unit interval across a sweep of rankings, a structural
166// invariant a regression in the formula would violate.
167func TestNDCGBoundedZeroOne(t *testing.T) {
168 relevant := rel("a", "b", "c")
169 rankings := [][]string{
170 {"a", "b", "c", "x", "y"},
171 {"x", "a", "y", "b", "c"},
172 {"y", "x", "c", "b", "a"},
173 {"a", "x", "b", "y", "c"},
174 {"x", "y", "z", "a", "b"},
175 }
176 for _, r := range rankings {
177 for k := 1; k <= 6; k++ {
178 got := NDCGAtK(r, relevant, k)
179 if got < 0.0 || got > 1.0+1e-9 {
180 t.Fatalf("NDCGAtK(%v, k=%d) = %v escaped [0,1]", r, k, got)
181 }
182 }
183 }
184}
185
186// TestCountRelevantIgnoresFalse proves a qrels map that records explicit non relevant judgments (value false)
187// does not inflate the relevant denominator, so recall and nDCG stay correct against such a fixture.
188func TestCountRelevantIgnoresFalse(t *testing.T) {
189 m := map[string]bool{"a": true, "b": false, "c": true}
190 // Two relevant (a, c). A ranking that finds only a is recall 0.5, not 0.333 from a count of three.
191 got := RecallAtK([]string{"a", "z"}, m, 5)
192 if !closeEnough(got, 0.5) {
193 t.Fatalf("RecallAtK with explicit false judgments = %v, want 0.5", got)
194 }
195}
196