Merge pull request #312 from mdekstrand/fix/309-main

Fix nDCG truncation bug (#309)
lenskit · Mar 11, 2022 · 371edfd · 371edfd
2 parents 0716300 + 65081ed
commit 371edfd
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/lenskit/metrics/topn.py b/lenskit/metrics/topn.py
@@ -300,11 +300,11 @@ def ndcg(recs, truth, discount=np.log2, k=None):
             The maximum list length.
     """
 
-    tpos = truth.index.get_indexer(recs['item'])
-
     if k is not None:
         recs = recs.iloc[:k]
 
+    tpos = truth.index.get_indexer(recs['item'])
+
     if 'rating' in truth.columns:
         i_rates = np.sort(truth.rating.values)[::-1]
         if k is not None:

diff --git a/tests/test_topn_ndcg.py b/tests/test_topn_ndcg.py
@@ -87,6 +87,14 @@ def test_ndcg_perfect():
     assert ndcg(recs, truth) == approx(1.0)
 
 
+def test_ndcg_perfect_k_short():
+    recs = pd.DataFrame({'item': [2, 3, 1]})
+    truth = pd.DataFrame({'item': [1, 2, 3], 'rating': [3.0, 5.0, 4.0]})
+    truth = truth.set_index('item')
+    assert ndcg(recs, truth, k=2) == approx(1.0)
+    assert ndcg(recs[:2], truth, k=2) == approx(1.0)
+
+
 def test_ndcg_wrong():
     recs = pd.DataFrame({'item': [1, 2]})
     truth = pd.DataFrame({'item': [1, 2, 3], 'rating': [3.0, 5.0, 4.0]})