Repository: ipums/hlink (language: python)

Issue:
Add F-measure to the computed model metrics, and include the raw confusion matrix in the output
F-measure is another helpful model metric, which can be computed in terms of precision and recall:

```
f-measure = 2 * ((precision * recall) / (precision + recall))
```

If you plug in the definitions of precision and recall in terms of true positives (`tp`), false positives (`fp`), and false negatives (`fn`), you get

```
f-measure = (2 * tp) / (2 * tp + fp + fn)
```

In addition to providing this metric, we should include the raw confusion matrix so that users can compute their own additional metrics if they would like to.


Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index 2adf10e..fd3991b 100644
--- a/.gitignore
+++ b/.gitignore
@@ -9,6 +9,7 @@ venv
 sphinx-docs/_*
 .coverage
 coverage_*
+.hypothesis/
 
 # Scala
 scala_jar/target
diff --git a/hlink/linking/core/model_metrics.py b/hlink/linking/core/model_metrics.py
new file mode 100644
index 0000000..d23fa00
--- /dev/null
+++ b/hlink/linking/core/model_metrics.py
@@ -0,0 +1,91 @@
+# This file is part of the ISRDI's hlink.
+# For copyright and licensing information, see the NOTICE and LICENSE files
+# in this project's top-level directory, and also on-line at:
+#   https://github.com/ipums/hlink
+"""
+Metrics for evaluating the performance of a machine learning model. These
+metrics operate on the "confusion matrix", which contains the four counts of
+true positives, true negatives, false positives, and false negatives.
+Throughout this module, we use the abbreviations true_pos, true_neg, false_pos,
+and false_neg for these confusion matrix values.
+
+All of these functions return math.nan in cases where they are not well-defined,
+such as cases with division by zero.
+"""
+
+import math
+
+
+def f_measure(true_pos: int, false_pos: int, false_neg: int) -> float:
+    """
+    Compute the F-measure, which is defined as the harmonic mean of precision
+    and recall:
+
+       2 * precision * recall / (precision + recall)
+
+    Using the definitions of precision and recall, we can write this in terms of
+    the confusion matrix entries as
+
+      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:
+    """
+    Compute the Matthews Correlation Coefficient (MCC). This can be written as
+    numerator / denominator, where
+
+      numerator = true_pos * true_neg - false_pos * false_neg
+
+    and
+
+      denominator = sqrt(
+        (true_pos + false_pos) *
+        (true_pos + false_neg) *
+        (true_neg + false_pos) *
+        (true_neg + false_neg)
+      )
+    """
+    denominator = math.sqrt(
+        (true_pos + false_pos)
+        * (true_pos + false_neg)
+        * (true_neg + false_pos)
+        * (true_neg + false_neg)
+    )
+    if denominator == 0:
+        return math.nan
+
+    numerator = true_pos * true_neg - false_pos * false_neg
+    return numerator / denominator
+
+
+def precision(true_pos: int, false_pos: int) -> float:
+    """
+    Compute the precision, also known as the positive predictive value (PPV).
+    This can be written in terms of the entries of the confusion matrix as
+
+      true_pos / (true_pos + false_pos)
+    """
+    denominator = true_pos + false_pos
+    if denominator == 0:
+        return math.nan
+
+    return true_pos / denominator
+
+
+def recall(true_pos: int, false_neg: int) -> float:
+    """
+    Compute the recall, which can be written in terms of the entries of the
+    confusion matrix as
+
+      true_pos / (true_pos + false_neg)
+    """
+    denominator = true_pos + false_neg
+    if denominator == 0:
+        return math.nan
+
+    return true_pos / denominator
diff --git a/pyproject.toml b/pyproject.toml
index 2a4b001..5c13c39 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -30,6 +30,7 @@ dependencies = [
 [project.optional-dependencies]
 dev = [
     "pytest>=7.1.0",
+    "hypothesis>=6.0",
     "black>=23.0",
     "flake8>=5.0",
     "pre-commit>=2.0",


Test patch:
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 4693b9a..63a1e3b 100644
--- a/hlink/linking/model_exploration/link_step_train_test_models.py
+++ b/hlink/linking/model_exploration/link_step_train_test_models.py
@@ -23,6 +23,7 @@
 from pyspark.sql import DataFrame
 from pyspark.sql.functions import col, count, count_if, mean
 from functools import reduce
+import hlink.linking.core.model_metrics as metrics_core
 import hlink.linking.core.threshold as threshold_core
 import hlink.linking.core.classifier as classifier_core
 
@@ -140,13 +141,18 @@ def make_threshold_matrix(self) -> list[list[float]]:
 # Both training and test results can be captured in this type
 @dataclass(kw_only=True)
 class ThresholdTestResult:
-    precision: float
-    recall: float
-    pr_auc: float
-    mcc: float
     model_id: str
     alpha_threshold: float
     threshold_ratio: float
+    true_pos: int
+    true_neg: int
+    false_pos: int
+    false_neg: int
+    precision: float
+    recall: float
+    mcc: float
+    f_measure: float
+    pr_auc: float
 
 
 class LinkStepTrainTestModels(LinkStep):
@@ -529,7 +535,7 @@ def _run(self) -> None:
         # threshold matrix entries.
         threshold_matrix_size = len(threshold_test_results[0])
 
-        thresholded_metrics_df = _create_thresholded_metrics_df()
+        thresholded_metrics_df = pd.DataFrame()
         for i in range(threshold_matrix_size):
             print(f"Aggregate threshold matrix entry {i}")
             thresholded_metrics_df = _aggregate_per_threshold_results(
@@ -538,13 +544,16 @@ def _run(self) -> None:
 
         print("***   Final thresholded metrics ***")
 
+        # Convert the parameters column to dtype string so that Spark can handle it
+        thresholded_metrics_df["parameters"] = thresholded_metrics_df[
+            "parameters"
+        ].apply(lambda t: str(t) if pd.notnull(t) else t)
         # thresholded_metrics_df has one row per threshold combination. and each outer fold
-        thresholded_metrics_df = _load_thresholded_metrics_df_params(
-            thresholded_metrics_df
-        )
-        _print_thresholded_metrics_df(
-            thresholded_metrics_df.sort_values(by="mcc_test_mean", ascending=False)
-        )
+        with pd.option_context(
+            "display.max_columns", None, "display.max_colwidth", None
+        ):
+            print(thresholded_metrics_df.sort_values(by="mcc_mean", ascending=False))
+        print("\n")
 
         self._save_training_results(thresholded_metrics_df, self.task.spark)
         self.task.spark.sql("set spark.sql.shuffle.partitions=200")
@@ -652,23 +661,29 @@ def _capture_prediction_results(
         predictions.createOrReplaceTempView(f"{table_prefix}predictions")
 
         (
-            tp_count,
-            fp_count,
-            fn_count,
-            tn_count,
+            true_pos,
+            false_pos,
+            false_neg,
+            true_neg,
         ) = _get_confusion_matrix(predictions, dep_var)
-        test_precision, test_recall, test_mcc = _get_aggregate_metrics(
-            tp_count, fp_count, fn_count, tn_count
-        )
+        precision = metrics_core.precision(true_pos, false_pos)
+        recall = metrics_core.recall(true_pos, false_neg)
+        mcc = metrics_core.mcc(true_pos, true_neg, false_pos, false_neg)
+        f_measure = metrics_core.f_measure(true_pos, false_pos, false_neg)
 
         result = ThresholdTestResult(
-            precision=test_precision,
-            recall=test_recall,
-            mcc=test_mcc,
-            pr_auc=pr_auc,
             model_id=model,
             alpha_threshold=alpha_threshold,
             threshold_ratio=threshold_ratio,
+            true_pos=true_pos,
+            true_neg=true_neg,
+            false_pos=false_pos,
+            false_neg=false_neg,
+            precision=precision,
+            recall=recall,
+            mcc=mcc,
+            f_measure=f_measure,
+            pr_auc=pr_auc,
         )
 
         return result
@@ -681,7 +696,6 @@ def _save_training_results(
         if desc_df.empty:
             print("Training results dataframe is empty.")
         else:
-            desc_df.dropna(axis=1, how="all", inplace=True)
             spark.createDataFrame(desc_df, samplingRatio=1).write.mode(
                 "overwrite"
             ).saveAsTable(f"{table_prefix}training_results")
@@ -690,21 +704,6 @@ def _save_training_results(
             # )
 
 
-def _calc_mcc(tp: int, tn: int, fp: int, fn: int) -> float:
-    """
-    Given the counts of true positives (tp), true negatives (tn), false
-    positives (fp), and false negatives (fn) for a model run, compute the
-    Matthews Correlation Coefficient (MCC).
-    """
-    if (math.sqrt((tp + fp) * (tp + fn) * (tn + fp) * (tn + fn))) != 0:
-        mcc = ((tp * tn) - (fp * fn)) / (
-            math.sqrt((tp + fp) * (tp + fn) * (tn + fp) * (tn + fn))
-        )
-    else:
-        mcc = 0
-    return mcc
-
-
 def _calc_threshold_matrix(
     alpha_threshold: float | list[float], threshold_ratio: float | list[float] | None
 ) -> list[list[float]]:
@@ -757,49 +756,26 @@ def _get_confusion_matrix(
     confusion matrix is the count of true positives, false positives, false
     negatives, and true negatives for the predictions.
 
-    Return a tuple (true_positives, false_positives, false_negatives,
-    true_negatives).
+    Return a tuple (true_pos, false_pos, false_neg, true_neg).
     """
     prediction_col = col("prediction")
     label_col = col(dep_var)
 
     confusion_matrix = predictions.select(
-        count_if((label_col == 1) & (prediction_col == 1)).alias("true_positives"),
-        count_if((label_col == 0) & (prediction_col == 1)).alias("false_positives"),
-        count_if((label_col == 1) & (prediction_col == 0)).alias("false_negatives"),
-        count_if((label_col == 0) & (prediction_col == 0)).alias("true_negatives"),
+        count_if((label_col == 1) & (prediction_col == 1)).alias("true_pos"),
+        count_if((label_col == 0) & (prediction_col == 1)).alias("false_pos"),
+        count_if((label_col == 1) & (prediction_col == 0)).alias("false_neg"),
+        count_if((label_col == 0) & (prediction_col == 0)).alias("true_neg"),
     )
     [confusion_row] = confusion_matrix.collect()
     return (
-        confusion_row.true_positives,
-        confusion_row.false_positives,
-        confusion_row.false_negatives,
-        confusion_row.true_negatives,
+        confusion_row.true_pos,
+        confusion_row.false_pos,
+        confusion_row.false_neg,
+        confusion_row.true_neg,
     )
 
 
-def _get_aggregate_metrics(
-    true_positives: int, false_positives: int, false_negatives: int, true_negatives: int
-) -> tuple[float, float, float]:
-    """
-    Given the counts of true positives, false positives, false negatives, and
-    true negatives for a model run, compute several metrics to evaluate the
-    model's quality.
-
-    Return a tuple of (precision, recall, Matthews Correlation Coefficient).
-    """
-    if (true_positives + false_positives) == 0:
-        precision = np.nan
-    else:
-        precision = true_positives / (true_positives + false_positives)
-    if (true_positives + false_negatives) == 0:
-        recall = np.nan
-    else:
-        recall = true_positives / (true_positives + false_negatives)
-    mcc = _calc_mcc(true_positives, true_negatives, false_positives, false_negatives)
-    return precision, recall, mcc
-
-
 # The outer list  entries hold results from each outer fold, the inner list has a ThresholdTestResult per threshold
 # matrix entry. We need to get data for each threshold entry together. Basically we need to invert the data.
 def _combine_by_threshold_matrix_entry(
@@ -831,6 +807,24 @@ def _combine_by_threshold_matrix_entry(
     return results
 
 
+def _compute_mean_and_stdev(values: list[float]) -> (float, float):
+    """
+    Given a list of floats, return a tuple (mean, stdev). If there aren't enough
+    values to compute the mean and/or stdev, return np.nan for that entry.
+    """
+    try:
+        mean = statistics.mean(values)
+    except statistics.StatisticsError:
+        mean = np.nan
+
+    try:
+        stdev = statistics.stdev(values)
+    except statistics.StatisticsError:
+        stdev = np.nan
+
+    return (mean, stdev)
+
+
 def _aggregate_per_threshold_results(
     thresholded_metrics_df: pd.DataFrame,
     prediction_results: list[ThresholdTestResult],
@@ -843,40 +837,17 @@ def _aggregate_per_threshold_results(
     threshold_ratio = prediction_results[0].threshold_ratio
 
     # Pull out columns to be aggregated
-    precision_test = [
-        r.precision for r in prediction_results if r.precision is not np.nan
-    ]
-    recall_test = [r.recall for r in prediction_results if r.recall is not np.nan]
-    pr_auc_test = [r.pr_auc for r in prediction_results if r.pr_auc is not np.nan]
-    mcc_test = [r.mcc for r in prediction_results if r.mcc is not np.nan]
-
-    # # variance requires at least two values
-    precision_test_sd = (
-        statistics.stdev(precision_test) if len(precision_test) > 1 else np.nan
-    )
-    recall_test_sd = statistics.stdev(recall_test) if len(recall_test) > 1 else np.nan
-    pr_auc_test_sd = statistics.stdev(pr_auc_test) if len(pr_auc_test) > 1 else np.nan
-    mcc_test_sd = statistics.stdev(mcc_test) if len(mcc_test) > 1 else np.nan
-
-    # Deal with tiny test data. This should never arise in practice but if it did we ought
-    # to issue a warning.
-    if len(precision_test) < 1:
-        #        raise RuntimeError("Not enough training data to get any valid precision values.")
-        precision_test_mean = np.nan
-    else:
-        precision_test_mean = (
-            statistics.mean(precision_test)
-            if len(precision_test) > 1
-            else precision_test[0]
-        )
-
-    if len(recall_test) < 1:
-        # raise RuntimeError("Not enough training data to get any valid recall values.")
-        recall_test_mean = np.nan
-    else:
-        recall_test_mean = (
-            statistics.mean(recall_test) if len(recall_test) > 1 else recall_test[0]
-        )
+    precision = [r.precision for r in prediction_results if not math.isnan(r.precision)]
+    recall = [r.recall for r in prediction_results if not math.isnan(r.recall)]
+    pr_auc = [r.pr_auc for r in prediction_results if not math.isnan(r.pr_auc)]
+    mcc = [r.mcc for r in prediction_results if not math.isnan(r.mcc)]
+    f_measure = [r.f_measure for r in prediction_results if not math.isnan(r.f_measure)]
+
+    (precision_mean, precision_sd) = _compute_mean_and_stdev(precision)
+    (recall_mean, recall_sd) = _compute_mean_and_stdev(recall)
+    (pr_auc_mean, pr_auc_sd) = _compute_mean_and_stdev(pr_auc)
+    (mcc_mean, mcc_sd) = _compute_mean_and_stdev(mcc)
+    (f_measure_mean, f_measure_sd) = _compute_mean_and_stdev(f_measure)
 
     new_desc = pd.DataFrame(
         {
@@ -884,14 +855,16 @@ def _aggregate_per_threshold_results(
             "parameters": [best_models[0].hyperparams],
             "alpha_threshold": [alpha_threshold],
             "threshold_ratio": [threshold_ratio],
-            "precision_test_mean": [precision_test_mean],
-            "precision_test_sd": [precision_test_sd],
-            "recall_test_mean": [recall_test_mean],
-            "recall_test_sd": [recall_test_sd],
-            "pr_auc_test_mean": [statistics.mean(pr_auc_test)],
-            "pr_auc_test_sd": [pr_auc_test_sd],
-            "mcc_test_mean": [statistics.mean(mcc_test)],
-            "mcc_test_sd": [mcc_test_sd],
+            "precision_mean": [precision_mean],
+            "precision_sd": [precision_sd],
+            "recall_mean": [recall_mean],
+            "recall_sd": [recall_sd],
+            "pr_auc_mean": [pr_auc_mean],
+            "pr_auc_sd": [pr_auc_sd],
+            "mcc_mean": [mcc_mean],
+            "mcc_sd": [mcc_sd],
+            "f_measure_mean": [f_measure_mean],
+            "f_measure_sd": [f_measure_sd],
         },
     )
 
@@ -902,57 +875,6 @@ def _aggregate_per_threshold_results(
     return thresholded_metrics_df
 
 
-def _print_thresholded_metrics_df(desc_df: pd.DataFrame) -> None:
-    pd.set_option("display.max_colwidth", None)
-    print(desc_df.iloc[-1])
-
-    print("\n")
-
-
-def _load_thresholded_metrics_df_params(desc_df: pd.DataFrame) -> pd.DataFrame:
-    params = [
-        "maxDepth",
-        "numTrees",
-        "featureSubsetStrategy",
-        "subsample",
-        "minInstancesPerNode",
-        "maxBins",
-        "class_weight",
-        "C",
-        "kernel",
-        "threshold",
-        "maxIter",
-    ]
-
-    load_params = lambda j, param: j.get(param, np.nan)
-    for param in params:
-        desc_df[param] = desc_df["parameters"].apply(load_params, args=(param,))
-    desc_df["class_weight"] = desc_df["class_weight"].apply(
-        lambda x: str(x) if pd.notnull(x) else x
-    )
-    desc_df["parameters"] = desc_df["parameters"].apply(
-        lambda t: str(t) if pd.notnull(t) else t
-    )
-    return desc_df
-
-
-def _create_thresholded_metrics_df() -> pd.DataFrame:
-    return pd.DataFrame(
-        columns=[
-            "model",
-            "parameters",
-            "alpha_threshold",
-            "threshold_ratio",
-            "precision_test_mean",
-            "precision_test_sd",
-            "recall_test_mean",
-            "recall_test_sd",
-            "mcc_test_mean",
-            "mcc_test_sd",
-        ]
-    )
-
-
 def _custom_param_grid_builder(
     model_parameters: list[dict[str, Any]]
 ) -> list[dict[str, Any]]:
diff --git a/hlink/tests/core/model_metrics_test.py b/hlink/tests/core/model_metrics_test.py
new file mode 100644
index 0000000..235ed75
--- /dev/null
+++ b/hlink/tests/core/model_metrics_test.py
@@ -0,0 +1,169 @@
+# This file is part of the ISRDI's hlink.
+# For copyright and licensing information, see the NOTICE and LICENSE files
+# in this project's top-level directory, and also on-line at:
+#   https://github.com/ipums/hlink
+import math
+
+from hypothesis import assume, given
+import hypothesis.strategies as st
+import pytest
+
+from hlink.linking.core.model_metrics import f_measure, mcc, precision, recall
+
+NonNegativeInt = st.integers(min_value=0)
+NegativeInt = st.integers(max_value=-1)
+
+
+def test_f_measure_example() -> None:
+    true_pos = 3112
+    false_pos = 205
+    false_neg = 1134
+
+    f_measure_score = f_measure(true_pos, false_pos, false_neg)
+    assert (
+        abs(f_measure_score - 0.8229539) < 0.0001
+    ), "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
+) -> None:
+    assume(true_pos + false_pos + false_neg > 0)
+    f_measure_score = f_measure(true_pos, false_pos, false_neg)
+    assert 0.0 <= f_measure_score <= 1.0
+
+
+@given(true_pos=NonNegativeInt, false_pos=NonNegativeInt, false_neg=NonNegativeInt)
+def test_f_measure_is_harmonic_mean_of_precision_and_recall(
+    true_pos: int, false_pos: int, false_neg: int
+) -> None:
+    precision_score = precision(true_pos, false_pos)
+    recall_score = recall(true_pos, false_neg)
+
+    assume(precision_score + recall_score > 0)
+
+    f_measure_score = f_measure(true_pos, false_pos, false_neg)
+    harmonic_mean = (
+        2 * precision_score * recall_score / (precision_score + recall_score)
+    )
+
+    assert (
+        abs(harmonic_mean - f_measure_score) < 0.0001
+    ), f"harmonic mean is {harmonic_mean}, but F-measure is {f_measure_score}"
+
+
+def test_mcc_example() -> None:
+    true_pos = 3112
+    false_pos = 205
+    false_neg = 1134
+    true_neg = 33259
+
+    mcc_score = mcc(true_pos, true_neg, false_pos, false_neg)
+    assert abs(mcc_score - 0.8111208) < 0.0001, "expected MCC to be near 0.8111208"
+
+
+@given(
+    true_pos=NonNegativeInt,
+    true_neg=NonNegativeInt,
+    false_pos=NonNegativeInt,
+    false_neg=NonNegativeInt,
+)
+def test_mcc_is_between_negative_1_and_positive_1(
+    true_pos: int, true_neg: int, false_pos: int, false_neg: int
+) -> None:
+    """
+    Under "normal circumstances", where the denominator of the Matthews Correlation
+    Coefficient isn't 0, its range is the interval [-1, 1].
+    """
+    assume(true_pos + false_pos > 0)
+    assume(true_pos + false_neg > 0)
+    assume(true_neg + false_pos > 0)
+    assume(true_neg + false_neg > 0)
+
+    mcc_score = mcc(true_pos, true_neg, false_pos, false_neg)
+    assert -1.0 <= mcc_score <= 1.0
+
+
+@pytest.mark.parametrize(
+    "true_pos,true_neg,false_pos,false_neg",
+    [(0, 0, 0, 0), (0, 1, 0, 1), (0, 1, 1, 0), (1, 0, 0, 1), (1, 0, 1, 0)],
+)
+def test_mcc_denom_zero(
+    true_pos: int, true_neg: int, false_pos: int, false_neg: int
+) -> None:
+    """
+    If the denominator of MCC is 0, it's not well-defined, and it returns NaN. This
+    can happen in a variety of situations if at least 2 of the inputs are 0.
+    """
+    mcc_score = mcc(true_pos, true_neg, false_pos, false_neg)
+    assert math.isnan(mcc_score)
+
+
+def test_precision_example() -> None:
+    true_pos = 3112
+    false_pos = 205
+
+    precision_score = precision(true_pos, false_pos)
+    assert (
+        abs(precision_score - 0.9381972) < 0.0001
+    ), "expected precision to be near 0.9381972"
+
+
+@given(true_pos=NonNegativeInt, false_pos=NonNegativeInt)
+def test_precision_between_0_and_1(true_pos: int, false_pos: int) -> None:
+    """
+    Under "normal circumstances" (there were at least some positive predictions)
+    precision()'s range is the interval [0.0, 1.0].
+    """
+    assume(true_pos + false_pos > 0)
+    precision_score = precision(true_pos, false_pos)
+    assert 0.0 <= precision_score <= 1.0
+
+
+def test_precision_no_positive_predictions() -> None:
+    """
+    When there are no positive predictions, true_pos=0 and false_pos=0, and
+    precision is not well defined. In this case we return NaN.
+    """
+    precision_score = precision(0, 0)
+    assert math.isnan(precision_score)
+
+
+def test_recall_example() -> None:
+    true_pos = 3112
+    false_neg = 1134
+
+    recall_score = recall(true_pos, false_neg)
+    assert (
+        abs(recall_score - 0.7329251) < 0.0001
+    ), "expected recall to be near 0.7329251"
+
+
+@given(true_pos=NonNegativeInt, false_neg=NonNegativeInt)
+def test_recall_between_0_and_1(true_pos: int, false_neg: int) -> None:
+    """
+    Under "normal circumstances" (there is at least one true positive or false
+    negative), the range of recall() is the interval [0.0, 1.0].
+    """
+    assume(true_pos + false_neg > 0)
+    recall_score = recall(true_pos, false_neg)
+    assert 0.0 <= recall_score <= 1.0
+
+
+def test_recall_no_true_pos_or_false_neg() -> None:
+    """
+    When both true_pos and false_neg are 0, recall is not well defined, and we
+    return NaN.
+    """
+    recall_score = recall(0, 0)
+    assert math.isnan(recall_score)
diff --git a/hlink/tests/hh_model_exploration_test.py b/hlink/tests/hh_model_exploration_test.py
index baa4d33..0e80026 100644
--- a/hlink/tests/hh_model_exploration_test.py
+++ b/hlink/tests/hh_model_exploration_test.py
@@ -54,10 +54,10 @@ def test_all_hh_mod_ev(
         "parameters",
         "alpha_threshold",
         "threshold_ratio",
-        "precision_test_mean",
-        "recall_test_mean",
-        "mcc_test_mean",
-        "pr_auc_test_mean",
+        "precision_mean",
+        "recall_mean",
+        "mcc_mean",
+        "pr_auc_mean",
     ]
 
     # TODO we should expect to get most of these columns once the results reporting is finished.
@@ -67,13 +67,13 @@ def test_all_hh_mod_ev(
         "alpha_threshold",
         "threshold_ratio",
         # "precision_test_mean",
-        "precision_test_sd",
-        "recall_test_mean",
-        "recall_test_sd",
-        "mcc_test_sd",
-        "mcc_test_mean",
-        "pr_auc_test_mean",
-        "pr_auc_test_sd",
+        "precision_sd",
+        "recall_mean",
+        "recall_sd",
+        "mcc_sd",
+        "mcc_mean",
+        "pr_auc_mean",
+        "pr_auc_sd",
         "maxDepth",
         "numTrees",
     ]
@@ -83,19 +83,15 @@ def test_all_hh_mod_ev(
 
     assert (
         0.6
-        < tr.query("model == 'logistic_regression'")["precision_test_mean"].iloc[0]
+        < tr.query("model == 'logistic_regression'")["precision_mean"].iloc[0]
         <= 1.0
     )
     assert tr.query("model == 'logistic_regression'")["alpha_threshold"].iloc[0] == 0.5
     assert (
-        0.7
-        < tr.query("model == 'logistic_regression'")["pr_auc_test_mean"].iloc[0]
-        <= 1.0
+        0.7 < tr.query("model == 'logistic_regression'")["pr_auc_mean"].iloc[0] <= 1.0
     )
     assert (
-        0.9
-        < tr.query("model == 'logistic_regression'")["recall_test_mean"].iloc[0]
-        <= 1.0
+        0.9 < tr.query("model == 'logistic_regression'")["recall_mean"].iloc[0] <= 1.0
     )
 
     preds = spark.table("hh_model_eval_predictions").toPandas()
diff --git a/hlink/tests/model_exploration_test.py b/hlink/tests/model_exploration_test.py
index 7222dbb..aad193e 100644
--- a/hlink/tests/model_exploration_test.py
+++ b/hlink/tests/model_exploration_test.py
@@ -13,7 +13,6 @@
     _custom_param_grid_builder,
     _get_model_parameters,
     _get_confusion_matrix,
-    _get_aggregate_metrics,
 )
 
 
@@ -696,8 +695,8 @@ def test_step_2_train_random_forest_spark(
     tr = spark.table("model_eval_training_results").toPandas()
     print(f"training results {tr}")
     # assert tr.shape == (1, 18)
-    assert tr.query("model == 'random_forest'")["pr_auc_test_mean"].iloc[0] > 2.0 / 3.0
-    assert tr.query("model == 'random_forest'")["maxDepth"].iloc[0] == 3
+    assert tr.query("model == 'random_forest'")["pr_auc_mean"].iloc[0] > 2.0 / 3.0
+    #  assert tr.query("model == 'random_forest'")["maxDepth"].iloc[0] == 3
 
     # TODO probably remove these since we're not planning to test suspicious data anymore.
     # I disabled the saving of suspicious in this test config so these are invalid currently.
@@ -732,10 +731,10 @@ def test_step_2_train_logistic_regression_spark(
     tr = spark.table("model_eval_training_results").toPandas()
     # assert tr.count == 3
 
-    assert tr.shape == (1, 11)
+    assert tr.shape == (1, 14)
     # This is now 0.83333333333.... I'm not sure it's worth testing against
     # assert tr.query("model == 'logistic_regression'")["pr_auc_mean"].iloc[0] == 0.75
-    assert tr.query("model == 'logistic_regression'")["pr_auc_test_mean"].iloc[0] > 0.74
+    assert tr.query("model == 'logistic_regression'")["pr_auc_mean"].iloc[0] > 0.74
     assert (
         round(tr.query("model == 'logistic_regression'")["alpha_threshold"].iloc[0], 1)
         == 0.7
@@ -760,12 +759,11 @@ def test_step_2_train_decision_tree_spark(
 
     print(f"Decision tree results: {tr}")
 
-    # TODO This is  1,12 instead of 1,13, because the precision_test_mean column is dropped as it is NaN
-    assert tr.shape == (1, 13)
-    # assert tr.query("model == 'decision_tree'")["precision_test_mean"].iloc[0] > 0
-    assert tr.query("model == 'decision_tree'")["maxDepth"].iloc[0] == 3
-    assert tr.query("model == 'decision_tree'")["minInstancesPerNode"].iloc[0] == 1
-    assert tr.query("model == 'decision_tree'")["maxBins"].iloc[0] == 7
+    assert tr.shape == (1, 14)
+    # assert tr.query("model == 'decision_tree'")["precision_mean"].iloc[0] > 0
+    #  assert tr.query("model == 'decision_tree'")["maxDepth"].iloc[0] == 3
+    #  assert tr.query("model == 'decision_tree'")["minInstancesPerNode"].iloc[0] == 1
+    #  assert tr.query("model == 'decision_tree'")["maxBins"].iloc[0] == 7
 
     main.do_drop_all("")
 
@@ -805,12 +803,12 @@ def test_step_2_train_gradient_boosted_trees_spark(
     # assert (
     #    tr.query("model == 'gradient_boosted_trees'")["precision_test_mean"].iloc[0] > 0
     # )
-    assert tr.query("model == 'gradient_boosted_trees'")["maxDepth"].iloc[0] == 5
-    assert (
-        tr.query("model == 'gradient_boosted_trees'")["minInstancesPerNode"].iloc[0]
-        == 1
-    )
-    assert tr.query("model == 'gradient_boosted_trees'")["maxBins"].iloc[0] == 5
+    #  assert tr.query("model == 'gradient_boosted_trees'")["maxDepth"].iloc[0] == 5
+    #  assert (
+    #  tr.query("model == 'gradient_boosted_trees'")["minInstancesPerNode"].iloc[0]
+    #  == 1
+    #  )
+    #  assert tr.query("model == 'gradient_boosted_trees'")["maxBins"].iloc[0] == 5
 
     main.do_drop_all("")
 
@@ -1016,20 +1014,3 @@ def test_get_confusion_matrix(spark: SparkSession) -> None:
     assert false_positives == 3
     assert false_negatives == 2
     assert true_negatives == 1
-
-
-def test_get_aggregate_metrics() -> None:
-    true_positives = 3112
-    false_positives = 205
-    false_negatives = 1134
-    true_negatives = 33259
-
-    precision, recall, mcc = _get_aggregate_metrics(
-        true_positives, false_positives, false_negatives, true_negatives
-    )
-
-    assert (
-        abs(precision - 0.9381972) < 0.0001
-    ), "expected precision to be near 0.9381972"
-    assert abs(recall - 0.7329251) < 0.0001, "expected recall to be near 0.7329251"
-    assert abs(mcc - 0.8111208) < 0.0001, "expected MCC to be near 0.8111208"