Implement NDCG cache. (#8893)

tests/cpp/common/test_ranking_utils.cc

@@ -1,16 +1,25 @@
 /**
  * Copyright 2023 by XGBoost Contributors
  */
-#include <gtest/gtest.h>                        // for Test, AssertionResult, Message, TestPartR...
-#include <gtest/gtest.h>                        // for ASSERT_NEAR, ASSERT_T...
-#include <xgboost/base.h>                       // for Args
+#include "test_ranking_utils.h"
+
+#include <gtest/gtest.h>
+#include <xgboost/base.h>                       // for Args, bst_group_t, kRtEps
 #include <xgboost/context.h>                    // for Context
+#include <xgboost/data.h>                       // for MetaInfo, DMatrix
+#include <xgboost/host_device_vector.h>         // for HostDeviceVector
+#include <xgboost/logging.h>                    // for Error
 #include <xgboost/string_view.h>                // for StringView
 
+#include <cstddef>                              // for size_t
 #include <cstdint>                              // for uint32_t
-#include <utility>                              // for pair
+#include <numeric>                              // for iota
+#include <utility>                              // for move
+#include <vector>                               // for vector
 
+#include "../../../src/common/numeric.h"        // for Iota
 #include "../../../src/common/ranking_utils.h"  // for LambdaRankParam, ParseMetricName, MakeMet...
+#include "../helpers.h"                         // for EmptyDMatrix
 
 namespace xgboost::ltr {
 TEST(RankingUtils, LambdaRankParam) {
@@ -66,4 +75,106 @@ TEST(RankingUtils, MakeMetricName) {
   name = MakeMetricName("map", 2, false);
   ASSERT_EQ(name, "map@2");
 }
+
+void TestRankingCache(Context const* ctx) {
+  auto p_fmat = EmptyDMatrix();
+  MetaInfo& info = p_fmat->Info();
+
+  info.num_row_ = 16;
+  info.labels.Reshape(info.num_row_);
+  auto& h_label = info.labels.Data()->HostVector();
+  for (std::size_t i = 0; i < h_label.size(); ++i) {
+    h_label[i] = i % 2;
+  }
+
+  LambdaRankParam param;
+  param.UpdateAllowUnknown(Args{});
+
+  RankingCache cache{ctx, info, param};
+
+  HostDeviceVector<float> predt(info.num_row_, 0);
+  auto& h_predt = predt.HostVector();
+  std::iota(h_predt.begin(), h_predt.end(), 0.0f);
+  predt.SetDevice(ctx->gpu_id);
+
+  auto rank_idx =
+      cache.SortedIdx(ctx, ctx->IsCPU() ? predt.ConstHostSpan() : predt.ConstDeviceSpan());
+
+  for (std::size_t i = 0; i < rank_idx.size(); ++i) {
+    ASSERT_EQ(rank_idx[i], rank_idx.size() - i - 1);
+  }
+}
+
+TEST(RankingCache, InitFromCPU) {
+  Context ctx;
+  TestRankingCache(&ctx);
+}
+
+void TestNDCGCache(Context const* ctx) {
+  auto p_fmat = EmptyDMatrix();
+  MetaInfo& info = p_fmat->Info();
+  LambdaRankParam param;
+  param.UpdateAllowUnknown(Args{});
+
+  {
+    // empty
+    NDCGCache cache{ctx, info, param};
+    ASSERT_EQ(cache.DataGroupPtr(ctx).size(), 2);
+  }
+
+  info.num_row_ = 3;
+  info.group_ptr_ = {static_cast<bst_group_t>(0), static_cast<bst_group_t>(info.num_row_)};
+
+  {
+    auto fail = [&]() { NDCGCache cache{ctx, info, param}; };
+    // empty label
+    ASSERT_THROW(fail(), dmlc::Error);
+    info.labels = linalg::Matrix<float>{{0.0f, 0.1f, 0.2f}, {3}, Context::kCpuId};
+    // invalid label
+    ASSERT_THROW(fail(), dmlc::Error);
+    auto h_labels = info.labels.HostView();
+    for (std::size_t i = 0; i < h_labels.Size(); ++i) {
+      h_labels(i) *= 10;
+    }
+    param.UpdateAllowUnknown(Args{{"ndcg_exp_gain", "false"}});
+    NDCGCache cache{ctx, info, param};
+    Context cpuctx;
+    auto inv_idcg = cache.InvIDCG(&cpuctx);
+    ASSERT_EQ(inv_idcg.Size(), 1);
+    ASSERT_NEAR(1.0 / inv_idcg(0), 2.63093, kRtEps);
+  }
+
+  {
+    param.UpdateAllowUnknown(Args{{"lambdarank_unbiased", "false"}});
+
+    std::vector<float> h_data(32);
+
+    common::Iota(ctx, h_data.begin(), h_data.end(), 0.0f);
+    info.labels.Reshape(h_data.size());
+    info.num_row_ = h_data.size();
+    info.group_ptr_.back() = info.num_row_;
+    info.labels.Data()->HostVector() = std::move(h_data);
+
+    {
+      NDCGCache cache{ctx, info, param};
+      Context cpuctx;
+      auto inv_idcg = cache.InvIDCG(&cpuctx);
+      ASSERT_NEAR(inv_idcg(0), 0.00551782, kRtEps);
+    }
+
+    param.UpdateAllowUnknown(
+        Args{{"lambdarank_num_pair_per_sample", "3"}, {"lambdarank_pair_method", "topk"}});
+    {
+      NDCGCache cache{ctx, info, param};
+      Context cpuctx;
+      auto inv_idcg = cache.InvIDCG(&cpuctx);
+      ASSERT_NEAR(inv_idcg(0), 0.01552123, kRtEps);
+    }
+  }
+}
+
+TEST(NDCGCache, InitFromCPU) {
+  Context ctx;
+  TestNDCGCache(&ctx);
+}
 }  // namespace xgboost::ltr