From 40f337ae094ee5f90a2d7fa48a6f8687f2fae783 Mon Sep 17 00:00:00 2001
From: Paul Sullivan <paulsullivan1@gmail.com>
Date: Wed, 11 Oct 2023 07:52:21 -0400
Subject: [PATCH 1/9] Discounted Cumulative Gain

---
 .../metrics/discounted_cumulative_gain.ex     | 86 +++++++++++++++++++
 .../discounted_cumulative_gain_test.exs       | 42 +++++++++
 2 files changed, 128 insertions(+)
 create mode 100644 lib/scholar/metrics/discounted_cumulative_gain.ex
 create mode 100644 test/scholar/metrics/discounted_cumulative_gain_test.exs

diff --git a/lib/scholar/metrics/discounted_cumulative_gain.ex b/lib/scholar/metrics/discounted_cumulative_gain.ex
new file mode 100644
index 00000000..5dbb747d
--- /dev/null
+++ b/lib/scholar/metrics/discounted_cumulative_gain.ex
@@ -0,0 +1,86 @@
+defmodule Scholar.Metrics.DiscountedCumulativeGain do
+  @moduledoc """
+  Discounted Cumulative Gain (DCG) is a measure of ranking quality.
+  It is based on the assumption that highly relevant documents appearing lower
+  in a search result list should be penalized as the graded relevance value is
+  reduced logarithmically proportional to the position of the result.
+  """
+
+  @doc """
+  Computes the DCG based on true relevancies (`y_true`) and predicted scores (`y_score`).
+  """
+
+  def compute(y_true, y_score, k \\ nil) do
+    # Ensure tensors are of the same shape
+    if Nx.shape(y_true) != Nx.shape(y_score) do
+      raise ArgumentError, "y_true and y_score tensors must have the same shape"
+    end
+
+    {adjusted_y_true, adjusted_y_score} = handle_ties(y_true, y_score)
+
+    sorted_indices = Nx.argsort(adjusted_y_score, axis: 0, direction: :desc)
+    sorted_y_true = Nx.take(adjusted_y_true, sorted_indices)
+
+    truncated_y_true = truncate_at_k(sorted_y_true, k)
+    dcg_value(truncated_y_true)
+  end
+  defp handle_ties(y_true, y_score) do
+    # Zip y_true and y_score together to work with pairs and convert to lists
+    zipped = y_true |> Nx.to_list() |> Enum.zip(Nx.to_list(y_score))
+
+    # Group items by their predicted scores and adjust groups if they contain ties
+    adjusted =
+      zipped
+      |> Enum.group_by(&elem(&1, 1))
+      |> Enum.flat_map(&adjust_group/1)
+
+    # Convert the lists back to tensors
+    {
+      Nx.tensor(Enum.map(adjusted, &elem(&1, 0))),
+      Nx.tensor(Enum.map(adjusted, &elem(&1, 1)))
+    }
+  end
+
+  # If a group has more than one element (i.e., there are ties), sort it by true_val
+  # and assign all elements the average rank. Otherwise, return the group unmodified.
+  defp adjust_group({_score, [single]}), do: [single]
+
+  defp adjust_group({score, group}) when length(group) > 1 do
+    group
+    |> Enum.sort_by(&elem(&1, 0), &>=/2)
+    |> Enum.map(&{elem(&1, 0), score})
+  end
+
+  defp dcg_value(y_true) do
+    float_y_true = Nx.as_type(y_true, :f32)
+
+    log_tensor =
+      y_true
+      |> Nx.shape()
+      |> Nx.iota()
+      |> Nx.as_type(:f32)
+      |> Nx.add(2.0)
+      |> Nx.log2()
+
+    if Enum.any?(Nx.to_flat_list(log_tensor), &(&1 < 0 or &1 !== &1)) do
+      raise ArithmeticError, "Encountered -Inf or NaN in log_tensor during DCG computation"
+    end
+
+    div_result = Nx.divide(float_y_true, log_tensor)
+
+    Nx.sum(div_result)
+  end
+
+  defp truncate_at_k(tensor, nil), do: tensor
+
+  defp truncate_at_k(tensor, k) do
+    shape = Nx.shape(tensor)
+
+    if Tuple.to_list(shape) |> Enum.at(0) > k do
+      {top_k, _rest} = Nx.split(tensor, k, axis: 0)
+      top_k
+    else
+      tensor
+    end
+  end
+end
diff --git a/test/scholar/metrics/discounted_cumulative_gain_test.exs b/test/scholar/metrics/discounted_cumulative_gain_test.exs
new file mode 100644
index 00000000..b4059537
--- /dev/null
+++ b/test/scholar/metrics/discounted_cumulative_gain_test.exs
@@ -0,0 +1,42 @@
+defmodule Scholar.Metrics.DiscountedCumulativeGainTest do
+  use Scholar.Case, async: true
+  alias Scholar.Metrics.DiscountedCumulativeGain
+
+  describe "compute/2" do
+    test "computes DCG when there are no ties" do
+      y_true = Nx.tensor([3, 2, 3, 0, 1, 2])
+      y_score = Nx.tensor([3.0, 2.2, 3.5, 0.5, 1.0, 2.1])
+
+      result = DiscountedCumulativeGain.compute(y_true, y_score)
+
+      assert %Nx.Tensor{data: data} = Nx.broadcast(result, {1})
+    end
+
+    test "computes DCG with ties" do
+      y_true = Nx.tensor([3, 3, 3])
+      y_score = Nx.tensor([2.0, 2.0, 3.5])
+
+      result = DiscountedCumulativeGain.compute(y_true, y_score)
+
+      assert %Nx.Tensor{data: data} = Nx.broadcast(result, {1})
+    end
+
+    test "raises error when shapes mismatch" do
+      y_true = Nx.tensor([3, 2, 3])
+      y_score = Nx.tensor([3.0, 2.2, 3.5, 0.5])
+
+      assert_raise ArgumentError, "y_true and y_score tensors must have the same shape", fn ->
+        DiscountedCumulativeGain.compute(y_true, y_score)
+      end
+    end
+
+    test "computes DCG for top-k values" do
+      y_true = Nx.tensor([3, 2, 3, 0, 1, 2])
+      y_score = Nx.tensor([3.0, 2.2, 3.5, 0.5, 1.0, 2.1])
+
+      result = DiscountedCumulativeGain.compute(y_true, y_score, 3)
+
+      assert %Nx.Tensor{data: data} = Nx.broadcast(result, {1})
+    end
+  end
+end

From d5e092c42a4d29da625d76a5655a66c1708bbb8e Mon Sep 17 00:00:00 2001
From: paulsullivanjr <paulsullivan1@gmail.com>
Date: Wed, 11 Oct 2023 13:00:05 -0400
Subject: [PATCH 2/9] Update lib/scholar/metrics/discounted_cumulative_gain.ex

Co-authored-by: Paulo Valente <16843419+polvalente@users.noreply.github.com>
---
 lib/scholar/metrics/discounted_cumulative_gain.ex | 1 -
 1 file changed, 1 deletion(-)

diff --git a/lib/scholar/metrics/discounted_cumulative_gain.ex b/lib/scholar/metrics/discounted_cumulative_gain.ex
index 5dbb747d..86f2409f 100644
--- a/lib/scholar/metrics/discounted_cumulative_gain.ex
+++ b/lib/scholar/metrics/discounted_cumulative_gain.ex
@@ -9,7 +9,6 @@ defmodule Scholar.Metrics.DiscountedCumulativeGain do
   @doc """
   Computes the DCG based on true relevancies (`y_true`) and predicted scores (`y_score`).
   """
-
   def compute(y_true, y_score, k \\ nil) do
     # Ensure tensors are of the same shape
     if Nx.shape(y_true) != Nx.shape(y_score) do

From 3f30d5a791b1add984126c4396a05dd416730a4a Mon Sep 17 00:00:00 2001
From: paulsullivanjr <paulsullivan1@gmail.com>
Date: Wed, 11 Oct 2023 13:00:38 -0400
Subject: [PATCH 3/9] Update lib/scholar/metrics/discounted_cumulative_gain.ex

Co-authored-by: Paulo Valente <16843419+polvalente@users.noreply.github.com>
---
 lib/scholar/metrics/discounted_cumulative_gain.ex | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/scholar/metrics/discounted_cumulative_gain.ex b/lib/scholar/metrics/discounted_cumulative_gain.ex
index 86f2409f..43066775 100644
--- a/lib/scholar/metrics/discounted_cumulative_gain.ex
+++ b/lib/scholar/metrics/discounted_cumulative_gain.ex
@@ -7,7 +7,7 @@ defmodule Scholar.Metrics.DiscountedCumulativeGain do
   """
 
   @doc """
-  Computes the DCG based on true relevancies (`y_true`) and predicted scores (`y_score`).
+  Computes the DCG based on true relevance scores (`y_true`) and their respective predicted scores (`y_score`).
   """
   def compute(y_true, y_score, k \\ nil) do
     # Ensure tensors are of the same shape

From 7611d3cd3ec119066c09b686a13bb75922faf7bd Mon Sep 17 00:00:00 2001
From: Paul Sullivan <paulsullivan1@gmail.com>
Date: Wed, 11 Oct 2023 21:22:29 -0400
Subject: [PATCH 4/9] Updated handle_ties to avoid the conversion

---
 .../metrics/discounted_cumulative_gain.ex     | 32 ++++++-------------
 1 file changed, 9 insertions(+), 23 deletions(-)

diff --git a/lib/scholar/metrics/discounted_cumulative_gain.ex b/lib/scholar/metrics/discounted_cumulative_gain.ex
index 43066775..fbf2f856 100644
--- a/lib/scholar/metrics/discounted_cumulative_gain.ex
+++ b/lib/scholar/metrics/discounted_cumulative_gain.ex
@@ -10,7 +10,6 @@ defmodule Scholar.Metrics.DiscountedCumulativeGain do
   Computes the DCG based on true relevance scores (`y_true`) and their respective predicted scores (`y_score`).
   """
   def compute(y_true, y_score, k \\ nil) do
-    # Ensure tensors are of the same shape
     if Nx.shape(y_true) != Nx.shape(y_score) do
       raise ArgumentError, "y_true and y_score tensors must have the same shape"
     end
@@ -23,31 +22,18 @@ defmodule Scholar.Metrics.DiscountedCumulativeGain do
     truncated_y_true = truncate_at_k(sorted_y_true, k)
     dcg_value(truncated_y_true)
   end
+
   defp handle_ties(y_true, y_score) do
-    # Zip y_true and y_score together to work with pairs and convert to lists
-    zipped = y_true |> Nx.to_list() |> Enum.zip(Nx.to_list(y_score))
-
-    # Group items by their predicted scores and adjust groups if they contain ties
-    adjusted =
-      zipped
-      |> Enum.group_by(&elem(&1, 1))
-      |> Enum.flat_map(&adjust_group/1)
-
-    # Convert the lists back to tensors
-    {
-      Nx.tensor(Enum.map(adjusted, &elem(&1, 0))),
-      Nx.tensor(Enum.map(adjusted, &elem(&1, 1)))
-    }
-  end
+    sorted_y_true = Nx.sort(y_true, axis: 0, direction: :desc)
+    sorted_y_score = Nx.sort(y_score, axis: 0, direction: :desc)
+
+    diff = Nx.diff(sorted_y_score)
+    selector = Nx.pad(diff, 1, [{1, 0, 0}])
+    adjusted_y_score = Nx.select(selector, sorted_y_score, 0)
 
-  # If a group has more than one element (i.e., there are ties), sort it by true_val
-  # and assign all elements the average rank. Otherwise, return the group unmodified.
-  defp adjust_group({_score, [single]}), do: [single]
+    adjusted_y_true = Nx.select(selector, sorted_y_true, 0)
 
-  defp adjust_group({score, group}) when length(group) > 1 do
-    group
-    |> Enum.sort_by(&elem(&1, 0), &>=/2)
-    |> Enum.map(&{elem(&1, 0), score})
+    {adjusted_y_true, adjusted_y_score}
   end
 
   defp dcg_value(y_true) do

From ad5b45ffd7d0dd7687aae8355f487ecdc6736ec4 Mon Sep 17 00:00:00 2001
From: Paul Sullivan <paulsullivan1@gmail.com>
Date: Tue, 17 Oct 2023 07:22:02 -0400
Subject: [PATCH 5/9] Updated to use defn and defnp, updated test

---
 .../metrics/discounted_cumulative_gain.ex     | 66 +++++++++++++------
 lib/scholar/options.ex                        |  3 +
 .../discounted_cumulative_gain_test.exs       |  8 ++-
 3 files changed, 53 insertions(+), 24 deletions(-)

diff --git a/lib/scholar/metrics/discounted_cumulative_gain.ex b/lib/scholar/metrics/discounted_cumulative_gain.ex
index fbf2f856..1522e13c 100644
--- a/lib/scholar/metrics/discounted_cumulative_gain.ex
+++ b/lib/scholar/metrics/discounted_cumulative_gain.ex
@@ -6,24 +6,50 @@ defmodule Scholar.Metrics.DiscountedCumulativeGain do
   reduced logarithmically proportional to the position of the result.
   """
 
+  import Nx.Defn
+  import Scholar.Shared
+  require Nx
+
+  opts = [
+    k: [
+      default: nil,
+      type: {:custom, Scholar.Options, :positive_number_or_nil, []},
+      doc: "Truncation parameter to consider only the top-k elements."
+    ]
+  ]
+
+  @opts_schema NimbleOptions.new!(opts)
+
+  deftransform compute(y_true, y_score, opts \\ []) do
+    compute_n(y_true, y_score, NimbleOptions.validate!(opts, @opts_schema))
+  end
+
   @doc """
+  ## Options
+  #{NimbleOptions.docs(@opts_schema)}
+
   Computes the DCG based on true relevance scores (`y_true`) and their respective predicted scores (`y_score`).
   """
-  def compute(y_true, y_score, k \\ nil) do
-    if Nx.shape(y_true) != Nx.shape(y_score) do
-      raise ArgumentError, "y_true and y_score tensors must have the same shape"
-    end
+  defn compute_n(y_true, y_score, opts) do
+    y_true_shape = Nx.shape(y_true)
+    y_score_shape = Nx.shape(y_score)
+
+    check_shape(y_true_shape, y_score_shape)
 
     {adjusted_y_true, adjusted_y_score} = handle_ties(y_true, y_score)
 
     sorted_indices = Nx.argsort(adjusted_y_score, axis: 0, direction: :desc)
     sorted_y_true = Nx.take(adjusted_y_true, sorted_indices)
 
-    truncated_y_true = truncate_at_k(sorted_y_true, k)
+    truncated_y_true = truncate_at_k(sorted_y_true, opts)
     dcg_value(truncated_y_true)
   end
 
-  defp handle_ties(y_true, y_score) do
+  defnp check_shape(y_true, y_pred) do
+    assert_same_shape!(y_true, y_pred)
+  end
+
+  defnp handle_ties(y_true, y_score) do
     sorted_y_true = Nx.sort(y_true, axis: 0, direction: :desc)
     sorted_y_score = Nx.sort(y_score, axis: 0, direction: :desc)
 
@@ -36,7 +62,7 @@ defmodule Scholar.Metrics.DiscountedCumulativeGain do
     {adjusted_y_true, adjusted_y_score}
   end
 
-  defp dcg_value(y_true) do
+  defnp dcg_value(y_true) do
     float_y_true = Nx.as_type(y_true, :f32)
 
     log_tensor =
@@ -47,25 +73,23 @@ defmodule Scholar.Metrics.DiscountedCumulativeGain do
       |> Nx.add(2.0)
       |> Nx.log2()
 
-    if Enum.any?(Nx.to_flat_list(log_tensor), &(&1 < 0 or &1 !== &1)) do
-      raise ArithmeticError, "Encountered -Inf or NaN in log_tensor during DCG computation"
-    end
-
     div_result = Nx.divide(float_y_true, log_tensor)
 
     Nx.sum(div_result)
   end
 
-  defp truncate_at_k(tensor, nil), do: tensor
-
-  defp truncate_at_k(tensor, k) do
-    shape = Nx.shape(tensor)
-
-    if Tuple.to_list(shape) |> Enum.at(0) > k do
-      {top_k, _rest} = Nx.split(tensor, k, axis: 0)
-      top_k
-    else
-      tensor
+  defnp truncate_at_k(tensor, opts) do
+    case opts[:k] do
+      nil ->
+        tensor
+
+      _ ->
+        if opts[:k] > Nx.axis_size(tensor, 0) do
+          tensor
+        else
+          {top_k, _rest} = Nx.split(tensor, opts[:k], axis: 0)
+          top_k
+        end
     end
   end
 end
diff --git a/lib/scholar/options.ex b/lib/scholar/options.ex
index a779d572..ffc3dafc 100644
--- a/lib/scholar/options.ex
+++ b/lib/scholar/options.ex
@@ -100,4 +100,7 @@ defmodule Scholar.Options do
     {:error,
      "expected metric to be a :cosine or tuple {:minkowski, p} where p is a positive number or :infinity, got: #{inspect(metric)}"}
   end
+
+  def positive_number_or_nil(nil), do: {:ok, nil}
+  def positive_number_or_nil(k), do: positive_number(k)
 end
diff --git a/test/scholar/metrics/discounted_cumulative_gain_test.exs b/test/scholar/metrics/discounted_cumulative_gain_test.exs
index b4059537..d6f19f7f 100644
--- a/test/scholar/metrics/discounted_cumulative_gain_test.exs
+++ b/test/scholar/metrics/discounted_cumulative_gain_test.exs
@@ -25,9 +25,11 @@ defmodule Scholar.Metrics.DiscountedCumulativeGainTest do
       y_true = Nx.tensor([3, 2, 3])
       y_score = Nx.tensor([3.0, 2.2, 3.5, 0.5])
 
-      assert_raise ArgumentError, "y_true and y_score tensors must have the same shape", fn ->
-        DiscountedCumulativeGain.compute(y_true, y_score)
-      end
+      assert_raise ArgumentError,
+                   "expected tensor to have shape {3}, got tensor with shape {4}",
+                   fn ->
+                     DiscountedCumulativeGain.compute(y_true, y_score)
+                   end
     end
 
     test "computes DCG for top-k values" do

From 3a2b9a9c8caa62f38fdef91077ec1c7f4c9d8953 Mon Sep 17 00:00:00 2001
From: Paul Sullivan <paulsullivan1@gmail.com>
Date: Sun, 22 Oct 2023 20:15:03 -0400
Subject: [PATCH 6/9] updated tests and handle_ties logic

---
 lib/scholar/metrics/discounted_cumulative_gain.ex    | 12 ++++++------
 .../metrics/discounted_cumulative_gain_test.exs      | 11 +++++++----
 2 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/lib/scholar/metrics/discounted_cumulative_gain.ex b/lib/scholar/metrics/discounted_cumulative_gain.ex
index 1522e13c..f4c73e06 100644
--- a/lib/scholar/metrics/discounted_cumulative_gain.ex
+++ b/lib/scholar/metrics/discounted_cumulative_gain.ex
@@ -50,14 +50,14 @@ defmodule Scholar.Metrics.DiscountedCumulativeGain do
   end
 
   defnp handle_ties(y_true, y_score) do
-    sorted_y_true = Nx.sort(y_true, axis: 0, direction: :desc)
-    sorted_y_score = Nx.sort(y_score, axis: 0, direction: :desc)
+    sorted_indices = Nx.argsort(y_score, axis: 0, direction: :desc)
 
-    diff = Nx.diff(sorted_y_score)
-    selector = Nx.pad(diff, 1, [{1, 0, 0}])
-    adjusted_y_score = Nx.select(selector, sorted_y_score, 0)
+    sorted_y_true = Nx.take(y_true, sorted_indices)
+    sorted_y_score = Nx.take(y_score, sorted_indices)
 
-    adjusted_y_true = Nx.select(selector, sorted_y_true, 0)
+    tie_sorted_indices = Nx.argsort(sorted_y_true, axis: 0, direction: :desc)
+    adjusted_y_true = Nx.take(sorted_y_true, tie_sorted_indices)
+    adjusted_y_score = Nx.take(sorted_y_score, tie_sorted_indices)
 
     {adjusted_y_true, adjusted_y_score}
   end
diff --git a/test/scholar/metrics/discounted_cumulative_gain_test.exs b/test/scholar/metrics/discounted_cumulative_gain_test.exs
index d6f19f7f..cd1df892 100644
--- a/test/scholar/metrics/discounted_cumulative_gain_test.exs
+++ b/test/scholar/metrics/discounted_cumulative_gain_test.exs
@@ -9,7 +9,8 @@ defmodule Scholar.Metrics.DiscountedCumulativeGainTest do
 
       result = DiscountedCumulativeGain.compute(y_true, y_score)
 
-      assert %Nx.Tensor{data: data} = Nx.broadcast(result, {1})
+      x = Nx.tensor([7.140995025634766])
+      assert x == Nx.broadcast(result, {1})
     end
 
     test "computes DCG with ties" do
@@ -18,7 +19,8 @@ defmodule Scholar.Metrics.DiscountedCumulativeGainTest do
 
       result = DiscountedCumulativeGain.compute(y_true, y_score)
 
-      assert %Nx.Tensor{data: data} = Nx.broadcast(result, {1})
+      x = Nx.tensor([6.3927892607143715])
+      assert x == Nx.broadcast(result, {1})
     end
 
     test "raises error when shapes mismatch" do
@@ -36,9 +38,10 @@ defmodule Scholar.Metrics.DiscountedCumulativeGainTest do
       y_true = Nx.tensor([3, 2, 3, 0, 1, 2])
       y_score = Nx.tensor([3.0, 2.2, 3.5, 0.5, 1.0, 2.1])
 
-      result = DiscountedCumulativeGain.compute(y_true, y_score, 3)
+      result = DiscountedCumulativeGain.compute(y_true, y_score, k: 3)
 
-      assert %Nx.Tensor{data: data} = Nx.broadcast(result, {1})
+      x = Nx.tensor([5.892789363861084])
+      assert x == Nx.broadcast(result, {1})
     end
   end
 end

From ce1e47a401ffbed16ea7727d7f18661308526d04 Mon Sep 17 00:00:00 2001
From: Paul Sullivan <paulsullivan1@gmail.com>
Date: Mon, 23 Oct 2023 07:21:03 -0400
Subject: [PATCH 7/9] Removed nil default for option

---
 lib/scholar/metrics/discounted_cumulative_gain.ex | 3 +--
 lib/scholar/options.ex                            | 3 ---
 2 files changed, 1 insertion(+), 5 deletions(-)

diff --git a/lib/scholar/metrics/discounted_cumulative_gain.ex b/lib/scholar/metrics/discounted_cumulative_gain.ex
index f4c73e06..9a630e7b 100644
--- a/lib/scholar/metrics/discounted_cumulative_gain.ex
+++ b/lib/scholar/metrics/discounted_cumulative_gain.ex
@@ -12,8 +12,7 @@ defmodule Scholar.Metrics.DiscountedCumulativeGain do
 
   opts = [
     k: [
-      default: nil,
-      type: {:custom, Scholar.Options, :positive_number_or_nil, []},
+      type: {:custom, Scholar.Options, :positive_number, []},
       doc: "Truncation parameter to consider only the top-k elements."
     ]
   ]
diff --git a/lib/scholar/options.ex b/lib/scholar/options.ex
index ffc3dafc..a779d572 100644
--- a/lib/scholar/options.ex
+++ b/lib/scholar/options.ex
@@ -100,7 +100,4 @@ defmodule Scholar.Options do
     {:error,
      "expected metric to be a :cosine or tuple {:minkowski, p} where p is a positive number or :infinity, got: #{inspect(metric)}"}
   end
-
-  def positive_number_or_nil(nil), do: {:ok, nil}
-  def positive_number_or_nil(k), do: positive_number(k)
 end

From 2cb9bcf8a4c2594a0291b9a869a43ceb11fdf5fc Mon Sep 17 00:00:00 2001
From: paulsullivanjr <paulsullivan1@gmail.com>
Date: Tue, 24 Oct 2023 06:50:18 -0400
Subject: [PATCH 8/9] Update lib/scholar/metrics/discounted_cumulative_gain.ex
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: José Valim <jose.valim@gmail.com>
---
 lib/scholar/metrics/discounted_cumulative_gain.ex | 1 +
 1 file changed, 1 insertion(+)

diff --git a/lib/scholar/metrics/discounted_cumulative_gain.ex b/lib/scholar/metrics/discounted_cumulative_gain.ex
index 9a630e7b..80421e3c 100644
--- a/lib/scholar/metrics/discounted_cumulative_gain.ex
+++ b/lib/scholar/metrics/discounted_cumulative_gain.ex
@@ -1,6 +1,7 @@
 defmodule Scholar.Metrics.DiscountedCumulativeGain do
   @moduledoc """
   Discounted Cumulative Gain (DCG) is a measure of ranking quality.
+
   It is based on the assumption that highly relevant documents appearing lower
   in a search result list should be penalized as the graded relevance value is
   reduced logarithmically proportional to the position of the result.

From 7cef63b30162cd0c318d7a812fab27821b986f91 Mon Sep 17 00:00:00 2001
From: Paul Sullivan <paulsullivan1@gmail.com>
Date: Fri, 3 Nov 2023 07:47:54 -0400
Subject: [PATCH 9/9] moved dcg into Scholar.Metrics.Ranking

---
 ...scounted_cumulative_gain.ex => ranking.ex} | 24 ++++++++++---------
 ...ulative_gain_test.exs => ranking_test.exs} | 14 +++++------
 2 files changed, 20 insertions(+), 18 deletions(-)
 rename lib/scholar/metrics/{discounted_cumulative_gain.ex => ranking.ex} (76%)
 rename test/scholar/metrics/{discounted_cumulative_gain_test.exs => ranking_test.exs} (72%)

diff --git a/lib/scholar/metrics/discounted_cumulative_gain.ex b/lib/scholar/metrics/ranking.ex
similarity index 76%
rename from lib/scholar/metrics/discounted_cumulative_gain.ex
rename to lib/scholar/metrics/ranking.ex
index 80421e3c..45a08447 100644
--- a/lib/scholar/metrics/discounted_cumulative_gain.ex
+++ b/lib/scholar/metrics/ranking.ex
@@ -1,36 +1,38 @@
-defmodule Scholar.Metrics.DiscountedCumulativeGain do
+defmodule Scholar.Metrics.Ranking do
   @moduledoc """
-  Discounted Cumulative Gain (DCG) is a measure of ranking quality.
+  Provides metrics and calculations related to ranking quality.
 
-  It is based on the assumption that highly relevant documents appearing lower
-  in a search result list should be penalized as the graded relevance value is
-  reduced logarithmically proportional to the position of the result.
+  Ranking metrics evaluate the quality of ordered lists of items,
+  often used in information retrieval and recommendation systems.
+
+  This module currently supports the following ranking metrics:
+    * Discounted Cumulative Gain (DCG)
   """
 
   import Nx.Defn
   import Scholar.Shared
   require Nx
 
-  opts = [
+  @dcg_opts [
     k: [
       type: {:custom, Scholar.Options, :positive_number, []},
       doc: "Truncation parameter to consider only the top-k elements."
     ]
   ]
 
-  @opts_schema NimbleOptions.new!(opts)
+  @dcg_opts_schema NimbleOptions.new!(@dcg_opts)
 
-  deftransform compute(y_true, y_score, opts \\ []) do
-    compute_n(y_true, y_score, NimbleOptions.validate!(opts, @opts_schema))
+  deftransform dcg(y_true, y_score, opts \\ []) do
+    dcg_n(y_true, y_score, NimbleOptions.validate!(opts, @dcg_opts_schema))
   end
 
   @doc """
   ## Options
-  #{NimbleOptions.docs(@opts_schema)}
+  #{NimbleOptions.docs(@dcg_opts_schema)}
 
   Computes the DCG based on true relevance scores (`y_true`) and their respective predicted scores (`y_score`).
   """
-  defn compute_n(y_true, y_score, opts) do
+  defn dcg_n(y_true, y_score, opts) do
     y_true_shape = Nx.shape(y_true)
     y_score_shape = Nx.shape(y_score)
 
diff --git a/test/scholar/metrics/discounted_cumulative_gain_test.exs b/test/scholar/metrics/ranking_test.exs
similarity index 72%
rename from test/scholar/metrics/discounted_cumulative_gain_test.exs
rename to test/scholar/metrics/ranking_test.exs
index cd1df892..a189ae52 100644
--- a/test/scholar/metrics/discounted_cumulative_gain_test.exs
+++ b/test/scholar/metrics/ranking_test.exs
@@ -1,13 +1,13 @@
-defmodule Scholar.Metrics.DiscountedCumulativeGainTest do
+defmodule Scholar.Metrics.RankingTest do
   use Scholar.Case, async: true
-  alias Scholar.Metrics.DiscountedCumulativeGain
+  alias Scholar.Metrics.Ranking
 
-  describe "compute/2" do
+  describe "dcg/3" do
     test "computes DCG when there are no ties" do
       y_true = Nx.tensor([3, 2, 3, 0, 1, 2])
       y_score = Nx.tensor([3.0, 2.2, 3.5, 0.5, 1.0, 2.1])
 
-      result = DiscountedCumulativeGain.compute(y_true, y_score)
+      result = Ranking.dcg(y_true, y_score)
 
       x = Nx.tensor([7.140995025634766])
       assert x == Nx.broadcast(result, {1})
@@ -17,7 +17,7 @@ defmodule Scholar.Metrics.DiscountedCumulativeGainTest do
       y_true = Nx.tensor([3, 3, 3])
       y_score = Nx.tensor([2.0, 2.0, 3.5])
 
-      result = DiscountedCumulativeGain.compute(y_true, y_score)
+      result = Ranking.dcg(y_true, y_score)
 
       x = Nx.tensor([6.3927892607143715])
       assert x == Nx.broadcast(result, {1})
@@ -30,7 +30,7 @@ defmodule Scholar.Metrics.DiscountedCumulativeGainTest do
       assert_raise ArgumentError,
                    "expected tensor to have shape {3}, got tensor with shape {4}",
                    fn ->
-                     DiscountedCumulativeGain.compute(y_true, y_score)
+                     Ranking.dcg(y_true, y_score)
                    end
     end
 
@@ -38,7 +38,7 @@ defmodule Scholar.Metrics.DiscountedCumulativeGainTest do
       y_true = Nx.tensor([3, 2, 3, 0, 1, 2])
       y_score = Nx.tensor([3.0, 2.2, 3.5, 0.5, 1.0, 2.1])
 
-      result = DiscountedCumulativeGain.compute(y_true, y_score, k: 3)
+      result = Ranking.dcg(y_true, y_score, k: 3)
 
       x = Nx.tensor([5.892789363861084])
       assert x == Nx.broadcast(result, {1})