From 26e99886ead8d265311fd439e16e476256198114 Mon Sep 17 00:00:00 2001
From: Daniel Tinoco <0urobor0s@users.noreply.github.com>
Date: Wed, 11 Oct 2023 15:03:50 +0100
Subject: [PATCH] F-beta score (#185)

---
 lib/scholar/metrics/classification.ex        | 163 +++++++++++++++----
 test/scholar/metrics/classification_test.exs |  20 +++
 2 files changed, 151 insertions(+), 32 deletions(-)

diff --git a/lib/scholar/metrics/classification.ex b/lib/scholar/metrics/classification.ex
index b2054757..d9bf2f7b 100644
--- a/lib/scholar/metrics/classification.ex
+++ b/lib/scholar/metrics/classification.ex
@@ -41,11 +41,13 @@ defmodule Scholar.Metrics.Classification do
           * `:micro` - Calculate metrics globally by counting the total true positives,
           false negatives and false positives.
 
-          * `:none` - The f1 scores for each class are returned.
+          * `:none` - The F-score values for each class are returned.
           """
         ]
       ]
 
+  fbeta_score_schema = f1_score_schema
+
   confusion_matrix_schema =
     general_schema ++
       [
@@ -163,6 +165,7 @@ defmodule Scholar.Metrics.Classification do
   @confusion_matrix_schema NimbleOptions.new!(confusion_matrix_schema)
   @balanced_accuracy_schema NimbleOptions.new!(balanced_accuracy_schema)
   @cohen_kappa_schema NimbleOptions.new!(cohen_kappa_schema)
+  @fbeta_score_schema NimbleOptions.new!(fbeta_score_schema)
   @f1_score_schema NimbleOptions.new!(f1_score_schema)
   @brier_score_loss_schema NimbleOptions.new!(brier_score_loss_schema)
   @accuracy_schema NimbleOptions.new!(accuracy_schema)
@@ -586,7 +589,7 @@ defmodule Scholar.Metrics.Classification do
   end
 
   @doc """
-  Calculates F1 score given rank-1 tensors which represent
+  Calculates F-beta score given rank-1 tensors which represent
   the expected (`y_true`) and predicted (`y_pred`) classes.
 
   If all examples are true negatives, then the result is 0 to
@@ -594,78 +597,174 @@ defmodule Scholar.Metrics.Classification do
 
   ## Options
 
-  #{NimbleOptions.docs(@f1_score_schema)}
+  #{NimbleOptions.docs(@fbeta_score_schema)}
 
   ## Examples
 
       iex> y_true = Nx.tensor([0, 1, 1, 1, 1, 0, 2, 1, 0, 1], type: :u32)
       iex> y_pred = Nx.tensor([0, 2, 1, 1, 2, 2, 2, 0, 0, 1], type: :u32)
-      iex> Scholar.Metrics.Classification.f1_score(y_true, y_pred, num_classes: 3)
+      iex> Scholar.Metrics.Classification.fbeta_score(y_true, y_pred, Nx.u32(1), num_classes: 3)
       #Nx.Tensor<
         f32[3]
         [0.6666666865348816, 0.6666666865348816, 0.4000000059604645]
       >
-      iex> Scholar.Metrics.Classification.f1_score(y_true, y_pred, num_classes: 3, average: :macro)
+      iex> Scholar.Metrics.Classification.fbeta_score(y_true, y_pred, Nx.u32(2), num_classes: 3)
+      #Nx.Tensor<
+        f32[3]
+        [0.6666666865348816, 0.5555555820465088, 0.625]
+      >
+      iex> Scholar.Metrics.Classification.fbeta_score(y_true, y_pred, Nx.f32(0.5), num_classes: 3)
+      #Nx.Tensor<
+        f32[3]
+        [0.6666666865348816, 0.8333333134651184, 0.29411765933036804]
+      >
+      iex> Scholar.Metrics.Classification.fbeta_score(y_true, y_pred, Nx.u32(2), num_classes: 3, average: :macro)
       #Nx.Tensor<
         f32
-        0.5777778029441833
+        0.6157407760620117
       >
-      iex> Scholar.Metrics.Classification.f1_score(y_true, y_pred, num_classes: 3, average: :weighted)
+      iex> Scholar.Metrics.Classification.fbeta_score(y_true, y_pred, Nx.u32(2), num_classes: 3, average: :weighted)
       #Nx.Tensor<
         f32
-        0.6399999856948853
+        0.5958333611488342
       >
-      iex> Scholar.Metrics.Classification.f1_score(y_true, y_pred, num_classes: 3, average: :micro)
+      iex> Scholar.Metrics.Classification.fbeta_score(y_true, y_pred, Nx.f32(0.5), num_classes: 3, average: :micro)
       #Nx.Tensor<
         f32
         0.6000000238418579
       >
-      iex> Scholar.Metrics.Classification.f1_score(Nx.tensor([1,0,1,0]), Nx.tensor([0, 1, 0, 1]), num_classes: 2, average: :none)
+      iex> Scholar.Metrics.Classification.fbeta_score(Nx.tensor([1, 0, 1, 0]), Nx.tensor([0, 1, 0, 1]), Nx.tensor(0.5), num_classes: 2, average: :none)
+      #Nx.Tensor<
+        f32[2]
+        [0.0, 0.0]
+      >
+      iex> Scholar.Metrics.Classification.fbeta_score(Nx.tensor([1, 0, 1, 0]), Nx.tensor([0, 1, 0, 1]), 0.5, num_classes: 2, average: :none)
       #Nx.Tensor<
         f32[2]
         [0.0, 0.0]
       >
   """
-  deftransform f1_score(y_true, y_pred, opts \\ []) do
-    f1_score_n(y_true, y_pred, NimbleOptions.validate!(opts, @f1_score_schema))
+  deftransform fbeta_score(y_true, y_pred, beta, opts \\ []) do
+    fbeta_score_n(y_true, y_pred, beta, NimbleOptions.validate!(opts, @fbeta_score_schema))
   end
 
-  defnp f1_score_n(y_true, y_pred, opts) do
+  defnp fbeta_score_n(y_true, y_pred, beta, opts) do
     check_shape(y_pred, y_true)
     num_classes = check_num_classes(opts[:num_classes])
+    average = opts[:average]
 
-    case opts[:average] do
+    {_precision, _recall, per_class_fscore} =
+      precision_recall_fscore_n(y_true, y_pred, beta, num_classes, average)
+
+    per_class_fscore
+  end
+
+  defnp fbeta_score_v(confusion_matrix, average) do
+    true_positive = Nx.take_diagonal(confusion_matrix)
+    false_positive = Nx.sum(confusion_matrix, axes: [0]) - true_positive
+    false_negative = Nx.sum(confusion_matrix, axes: [1]) - true_positive
+
+    case average do
       :micro ->
-        accuracy(y_true, y_pred)
+        true_positive = Nx.sum(true_positive)
+        false_positive = Nx.sum(false_positive)
+        false_negative = Nx.sum(false_negative)
+
+        {true_positive, false_positive, false_negative}
 
       _ ->
-        cm = confusion_matrix(y_true, y_pred, num_classes: num_classes)
-        true_positive = Nx.take_diagonal(cm)
-        false_positive = Nx.sum(cm, axes: [0]) - true_positive
-        false_negative = Nx.sum(cm, axes: [1]) - true_positive
+        {true_positive, false_positive, false_negative}
+    end
+  end
 
-        precision = safe_division(true_positive, true_positive + false_positive)
+  defnp precision_recall_fscore_n(y_true, y_pred, beta, num_classes, average) do
+    confusion_matrix = confusion_matrix(y_true, y_pred, num_classes: num_classes)
+    {true_positive, false_positive, false_negative} = fbeta_score_v(confusion_matrix, average)
 
-        recall = safe_division(true_positive, true_positive + false_negative)
+    precision = safe_division(true_positive, true_positive + false_positive)
+    recall = safe_division(true_positive, true_positive + false_negative)
 
-        per_class_f1 = safe_division(2 * precision * recall, precision + recall)
+    per_class_fscore =
+      cond do
+        # Should only be +Inf
+        Nx.is_infinity(beta) ->
+          recall
 
-        case opts[:average] do
-          :none ->
-            per_class_f1
+        beta == 0 ->
+          precision
 
-          :macro ->
-            Nx.mean(per_class_f1)
+        true ->
+          beta2 = Nx.pow(beta, 2)
+          safe_division((1 + beta2) * precision * recall, beta2 * precision + recall)
+      end
 
-          :weighted ->
-            support = (y_true == Nx.iota({num_classes, 1})) |> Nx.sum(axes: [1])
+    case average do
+      :none ->
+        {precision, recall, per_class_fscore}
 
-            safe_division(per_class_f1 * support, Nx.sum(support))
-            |> Nx.sum()
-        end
+      :micro ->
+        {precision, recall, per_class_fscore}
+
+      :macro ->
+        {precision, recall, Nx.mean(per_class_fscore)}
+
+      :weighted ->
+        support = (y_true == Nx.iota({num_classes, 1})) |> Nx.sum(axes: [1])
+
+        per_class_fscore =
+          (per_class_fscore * support)
+          |> safe_division(Nx.sum(support))
+          |> Nx.sum()
+
+        {precision, recall, per_class_fscore}
     end
   end
 
+  @doc """
+  Calculates F1 score given rank-1 tensors which represent
+  the expected (`y_true`) and predicted (`y_pred`) classes.
+
+  If all examples are true negatives, then the result is 0 to
+  avoid zero division.
+
+  ## Options
+
+  #{NimbleOptions.docs(@f1_score_schema)}
+
+  ## Examples
+
+      iex> y_true = Nx.tensor([0, 1, 1, 1, 1, 0, 2, 1, 0, 1], type: :u32)
+      iex> y_pred = Nx.tensor([0, 2, 1, 1, 2, 2, 2, 0, 0, 1], type: :u32)
+      iex> Scholar.Metrics.Classification.f1_score(y_true, y_pred, num_classes: 3)
+      #Nx.Tensor<
+        f32[3]
+        [0.6666666865348816, 0.6666666865348816, 0.4000000059604645]
+      >
+      iex> Scholar.Metrics.Classification.f1_score(y_true, y_pred, num_classes: 3, average: :macro)
+      #Nx.Tensor<
+        f32
+        0.5777778029441833
+      >
+      iex> Scholar.Metrics.Classification.f1_score(y_true, y_pred, num_classes: 3, average: :weighted)
+      #Nx.Tensor<
+        f32
+        0.6399999856948853
+      >
+      iex> Scholar.Metrics.Classification.f1_score(y_true, y_pred, num_classes: 3, average: :micro)
+      #Nx.Tensor<
+        f32
+        0.6000000238418579
+      >
+      iex> Scholar.Metrics.Classification.f1_score(Nx.tensor([1, 0, 1, 0]), Nx.tensor([0, 1, 0, 1]), num_classes: 2, average: :none)
+      #Nx.Tensor<
+        f32[2]
+        [0.0, 0.0]
+      >
+  """
+  deftransform f1_score(y_true, y_pred, opts \\ []) do
+    fbeta_score_n(y_true, y_pred, 1, NimbleOptions.validate!(opts, @f1_score_schema))
+  end
+
   @doc """
   Zero-one classification loss.
 
diff --git a/test/scholar/metrics/classification_test.exs b/test/scholar/metrics/classification_test.exs
index ed00802f..7667e2df 100644
--- a/test/scholar/metrics/classification_test.exs
+++ b/test/scholar/metrics/classification_test.exs
@@ -14,4 +14,24 @@ defmodule Scholar.Metrics.ClassificationTest do
     assert_all_close(tpr, Nx.tensor([0.0, 0.5, 1.0, 1.0]))
     assert_all_close(thresholds, Nx.tensor([1.3, 0.3, 0.2, 0.1]))
   end
+
+  describe "fbeta_score" do
+    test "equals recall when beta is infinity" do
+      beta = Nx.tensor(:infinity)
+      y_true = Nx.tensor([0, 0, 0, 0, 0, 1, 1, 1, 1, 1], type: :u32)
+      y_pred = Nx.tensor([1, 1, 1, 1, 1, 1, 1, 1, 1, 1], type: :u32)
+      fbeta_scores = Classification.fbeta_score(y_true, y_pred, beta, num_classes: 2)
+
+      assert_all_close(fbeta_scores, Classification.recall(y_true, y_pred, num_classes: 2))
+    end
+
+    test "equals precision when beta is 0" do
+      beta = 0
+      y_true = Nx.tensor([0, 0, 0, 0, 0, 1, 1, 1, 1, 1], type: :u32)
+      y_pred = Nx.tensor([1, 1, 1, 1, 1, 1, 1, 1, 1, 1], type: :u32)
+      fbeta_scores = Classification.fbeta_score(y_true, y_pred, beta, num_classes: 2)
+
+      assert_all_close(fbeta_scores, Classification.precision(y_true, y_pred, num_classes: 2))
+    end
+  end
 end