From 7f0c48c64769b1bcaf81ca72b8102497f972e509 Mon Sep 17 00:00:00 2001
From: rileyh <rileyh@umn.edu>
Date: Thu, 12 Dec 2024 09:40:34 -0600
Subject: [PATCH] [#179] Include F-measure in ThresholdTestResults

This required fixing a bug in core.model_metrics.f_measure where it errored out
instead of returning NaN when its denominator was 0.
---
 hlink/linking/core/model_metrics.py                      | 5 ++++-
 .../model_exploration/link_step_train_test_models.py     | 5 ++++-
 hlink/tests/core/model_metrics_test.py                   | 9 +++++++++
 3 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/hlink/linking/core/model_metrics.py b/hlink/linking/core/model_metrics.py
index 769feee..95b5ef8 100644
--- a/hlink/linking/core/model_metrics.py
+++ b/hlink/linking/core/model_metrics.py
@@ -6,7 +6,10 @@
 
 
 def f_measure(true_pos: int, false_pos: int, false_neg: int) -> float:
-    return 2 * true_pos / (2 * true_pos + false_pos + false_neg)
+    denominator = 2 * true_pos + false_pos + false_neg
+    if denominator == 0:
+        return math.nan
+    return 2 * true_pos / denominator
 
 
 def mcc(true_pos: int, true_neg: int, false_pos: int, false_neg: int) -> float:
diff --git a/hlink/linking/model_exploration/link_step_train_test_models.py b/hlink/linking/model_exploration/link_step_train_test_models.py
index 31f1d63..78f7f21 100644
--- a/hlink/linking/model_exploration/link_step_train_test_models.py
+++ b/hlink/linking/model_exploration/link_step_train_test_models.py
@@ -143,8 +143,9 @@ def make_threshold_matrix(self) -> list[list[float]]:
 class ThresholdTestResult:
     precision: float
     recall: float
-    pr_auc: float
     mcc: float
+    f_measure: float
+    pr_auc: float
     model_id: str
     alpha_threshold: float
     threshold_ratio: float
@@ -661,11 +662,13 @@ def _capture_prediction_results(
         precision = metrics_core.precision(tp_count, fp_count)
         recall = metrics_core.recall(tp_count, fn_count)
         mcc = metrics_core.mcc(tp_count, tn_count, fp_count, fn_count)
+        f_measure = metrics_core.f_measure(tp_count, fp_count, fn_count)
 
         result = ThresholdTestResult(
             precision=precision,
             recall=recall,
             mcc=mcc,
+            f_measure=f_measure,
             pr_auc=pr_auc,
             model_id=model,
             alpha_threshold=alpha_threshold,
diff --git a/hlink/tests/core/model_metrics_test.py b/hlink/tests/core/model_metrics_test.py
index 2cb1d33..56df30c 100644
--- a/hlink/tests/core/model_metrics_test.py
+++ b/hlink/tests/core/model_metrics_test.py
@@ -24,6 +24,15 @@ def test_f_measure_example() -> None:
     ), "expected F-measure to be near 0.8229539"
 
 
+def test_f_measure_all_zeroes() -> None:
+    """
+    When true_pos, false_pos, and false_neg are all 0, f_measure is undefined and
+    returns NaN to indicate this.
+    """
+    f_measure_score = f_measure(0, 0, 0)
+    assert math.isnan(f_measure_score)
+
+
 @given(true_pos=NonNegativeInt, false_pos=NonNegativeInt, false_neg=NonNegativeInt)
 def test_f_measure_between_0_and_1(
     true_pos: int, false_pos: int, false_neg: int