Repository: swisstopo/swissgeol-boreholes-dataextraction

Issue:
Find a better way to associate the "correctness" from the evaluation with the extracted groundwater values
[Migrated from LGVISIUM-65.]

This issue originated from the PR of the groundwater level: https://github.com/swisstopo/swissgeol-boreholes-dataextraction/pull/72#issuecomment-2314949349 

The goals is to find a better way to associate the "correctness" from the evaluation with the individual extracted values, so that we can show (red/green) the correctness of each extracted groundwater measurement also in the visualisation.

In particular, solve the TODOs in `groundwater_evaluator.py` and `draw.py` that were added in https://github.com/swisstopo/swissgeol-boreholes-dataextraction/pull/148.

Discussion:

Solution patch:
diff --git a/src/stratigraphy/annotations/draw.py b/src/stratigraphy/annotations/draw.py
index a6c65147..f1f03700 100644
--- a/src/stratigraphy/annotations/draw.py
+++ b/src/stratigraphy/annotations/draw.py
@@ -81,9 +81,12 @@ def draw_predictions(
                             )
                         for groundwater_entry in groundwaters.groundwater_feature_list:
                             if page_number == groundwater_entry.page:
-                                # TODO: store correctness with the groundwater object
-                                # See also https://github.com/swisstopo/swissgeol-boreholes-dataextraction/issues/124
-                                draw_feature(shape, groundwater_entry.rect * page.derotation_matrix, None, "pink")
+                                draw_feature(
+                                    shape,
+                                    groundwater_entry.rect * page.derotation_matrix,
+                                    groundwater_entry.feature.is_correct,
+                                    "pink",
+                                )
                         draw_depth_columns_and_material_rect(
                             shape,
                             page.derotation_matrix,
diff --git a/src/stratigraphy/evaluation/groundwater_evaluator.py b/src/stratigraphy/evaluation/groundwater_evaluator.py
index 2ada6deb..09607beb 100644
--- a/src/stratigraphy/evaluation/groundwater_evaluator.py
+++ b/src/stratigraphy/evaluation/groundwater_evaluator.py
@@ -4,7 +4,7 @@
 
 from stratigraphy.benchmark.metrics import OverallMetrics
 from stratigraphy.evaluation.evaluation_dataclasses import Metrics
-from stratigraphy.evaluation.utility import count_against_ground_truth
+from stratigraphy.evaluation.utility import evaluate
 from stratigraphy.groundwater.groundwater_extraction import Groundwater
 from stratigraphy.util.borehole_predictions import FileGroundwaterWithGroundTruth
 
@@ -92,35 +92,26 @@ def evaluate(self) -> OverallGroundwaterMetrics:
                     for json_gt_data in borehole_data.ground_truth
                 ]
 
-                # TODO store the correctness directly on the Groundwater objects, so we can use that in the
-                # visualizations (cf. https://github.com/swisstopo/swissgeol-boreholes-dataextraction/issues/124)
-                entries = borehole_data.groundwater.groundwater_feature_list if borehole_data.groundwater else []
-                groundwater_metrics = count_against_ground_truth(
-                    [
-                        (
-                            entry.feature.depth,
-                            entry.feature.format_date(),
-                            entry.feature.elevation,
-                        )
-                        for entry in entries
-                    ],
-                    [(entry.depth, entry.format_date(), entry.elevation) for entry in gt_groundwater],
+                entries = (
+                    [feature_on_page.feature for feature_on_page in borehole_data.groundwater.groundwater_feature_list]
+                    if borehole_data.groundwater
+                    else []
                 )
-                groundwater_depth_metrics = count_against_ground_truth(
-                    [entry.feature.depth for entry in entries],
-                    [entry.depth for entry in gt_groundwater],
-                )
-                groundwater_elevation_metrics = count_against_ground_truth(
-                    [entry.feature.elevation for entry in entries],
-                    [entry.elevation for entry in gt_groundwater],
-                )
-                groundwater_date_metrics = count_against_ground_truth(
-                    [entry.feature.format_date() for entry in entries],
-                    [entry.format_date() for entry in gt_groundwater],
-                )
-                groundwater_metrics_list.append(groundwater_metrics)
+
+                groundwater_evaluation = evaluate(entries, gt_groundwater, self.match_groundwater)
+                groundwater_metrics_list.append(groundwater_evaluation.metrics)
+                for feature, is_correct in zip(entries, groundwater_evaluation.extracted_correct, strict=True):
+                    feature.is_correct = is_correct
+
+                groundwater_depth_metrics = evaluate(entries, gt_groundwater, self.match_groundwater_depth).metrics
                 groundwater_depth_metrics_list.append(groundwater_depth_metrics)
+
+                groundwater_elevation_metrics = evaluate(
+                    entries, gt_groundwater, self.match_groundwater_elevation
+                ).metrics
                 groundwater_elevation_metrics_list.append(groundwater_elevation_metrics)
+
+                groundwater_date_metrics = evaluate(entries, gt_groundwater, self.match_groundwater_date).metrics
                 groundwater_date_metrics_list.append(groundwater_date_metrics)
 
             # we take the micro-average across boreholes
@@ -135,3 +126,21 @@ def evaluate(self) -> OverallGroundwaterMetrics:
             overall_groundwater_metrics.add_groundwater_metrics(file_groundwater_metrics)
 
         return overall_groundwater_metrics
+
+    @staticmethod
+    def match_groundwater(extracted: Groundwater, ground_truth: Groundwater) -> bool:
+        extracted_tuple = (extracted.depth, extracted.format_date(), extracted.elevation)
+        ground_truth_tuple = (ground_truth.depth, ground_truth.format_date(), ground_truth.elevation)
+        return extracted_tuple == ground_truth_tuple
+
+    @staticmethod
+    def match_groundwater_depth(extracted: Groundwater, ground_truth: Groundwater) -> bool:
+        return extracted.depth == ground_truth.depth
+
+    @staticmethod
+    def match_groundwater_date(extracted: Groundwater, ground_truth: Groundwater) -> bool:
+        return extracted.format_date() == ground_truth.format_date()
+
+    @staticmethod
+    def match_groundwater_elevation(extracted: Groundwater, ground_truth: Groundwater) -> bool:
+        return extracted.elevation == ground_truth.elevation
diff --git a/src/stratigraphy/evaluation/metadata_evaluator.py b/src/stratigraphy/evaluation/metadata_evaluator.py
index 62c8a7ad..f1ffda12 100644
--- a/src/stratigraphy/evaluation/metadata_evaluator.py
+++ b/src/stratigraphy/evaluation/metadata_evaluator.py
@@ -7,6 +7,7 @@
     Metrics,
     OverallBoreholeMetadataMetrics,
 )
+from stratigraphy.evaluation.utility import evaluate_single
 from stratigraphy.metadata.coordinate_extraction import Coordinate
 from stratigraphy.util.borehole_predictions import FileMetadataWithGroundTruth
 
@@ -49,7 +50,10 @@ def evaluate(self) -> OverallBoreholeMetadataMetrics:
                 )
                 ground_truth_coordinates = borehole_data.ground_truth.get("coordinates")
 
-                coordinate_metrics = self._evaluate_coordinate(extracted_coordinates, ground_truth_coordinates)
+                evaluation_result = evaluate_single(
+                    extracted_coordinates, ground_truth_coordinates, self.match_coordinates
+                )
+                coordinate_metrics = evaluation_result.metrics
                 if borehole_data.metadata and borehole_data.metadata.coordinates:
                     borehole_data.metadata.coordinates.feature.is_correct = coordinate_metrics.tp > 0
                 coordinate_metrics_list.append(coordinate_metrics)
@@ -63,7 +67,8 @@ def evaluate(self) -> OverallBoreholeMetadataMetrics:
                     else None
                 )
                 ground_truth_elevation = borehole_data.ground_truth.get("reference_elevation")
-                elevation_metrics = self._evaluate_elevation(extracted_elevation, ground_truth_elevation)
+                evaluation_result = evaluate_single(extracted_elevation, ground_truth_elevation, self.match_elevation)
+                elevation_metrics = evaluation_result.metrics
                 if borehole_data.metadata and borehole_data.metadata.elevation:
                     borehole_data.metadata.elevation.feature.is_correct = elevation_metrics.tp > 0
                 elevation_metrics_list.append(elevation_metrics)
@@ -79,78 +84,40 @@ def evaluate(self) -> OverallBoreholeMetadataMetrics:
 
         return metadata_metrics_list
 
-    def _evaluate_elevation(self, extracted_elevation: float | None, ground_truth_elevation: float | None):
-        """Private method used to evaluate the extracted elevation against the ground truth.
+    @staticmethod
+    def match_elevation(extracted_elevation: float, ground_truth_elevation: float):
+        """Method used to evaluate the extracted elevation against the ground truth.
 
         Args:
-            extracted_elevation (float | None): the extracted elevation
-            ground_truth_elevation (float | None): the groundtruth elevation
+            extracted_elevation (float): the extracted elevation
+            ground_truth_elevation (float): the groundtruth elevation
 
         Returns:
-            Metrics: the metric for this elevation.
+            bool: if the extracted evaluation matches the ground truth
         """
-        if extracted_elevation is not None and ground_truth_elevation is not None:
-            if math.isclose(extracted_elevation, ground_truth_elevation, abs_tol=0.1):
-                elevation_metrics = Metrics(
-                    tp=1,
-                    fp=0,
-                    fn=0,
-                )
-            else:
-                elevation_metrics = Metrics(
-                    tp=0,
-                    fp=1,
-                    fn=1,
-                )
-        else:
-            elevation_metrics = Metrics(
-                tp=0,
-                fp=1 if extracted_elevation is not None else 0,
-                fn=1 if ground_truth_elevation is not None else 0,
-            )
+        return math.isclose(extracted_elevation, ground_truth_elevation, abs_tol=0.1)
 
-        return elevation_metrics
-
-    def _evaluate_coordinate(self, extracted_coordinates: Coordinate | None, ground_truth_coordinates: dict | None):
-        """Private method used to evaluate the extracted coordinates against the ground truth.
+    @staticmethod
+    def match_coordinates(extracted_coordinates: Coordinate, ground_truth_coordinates: dict):
+        """Method used to evaluate the extracted coordinates against the ground truth.
 
         Args:
-            extracted_coordinates (Coordinate | None): the extracted coordinates
-            ground_truth_coordinates (dict | None): the groundtruth coordinates
+            extracted_coordinates (Coordinate): the extracted coordinates
+            ground_truth_coordinates (dict): the groundtruth coordinates
 
         Returns:
-            Metrics: the metric for these coordinates.
+            bool: if the extracted cooredinates match the ground truth
         """
-        if extracted_coordinates and ground_truth_coordinates:
-            if extracted_coordinates.east.coordinate_value > 2e6 and ground_truth_coordinates["E"] < 2e6:
-                ground_truth_east = int(ground_truth_coordinates["E"]) + 2e6
-                ground_truth_north = int(ground_truth_coordinates["N"]) + 1e6
-            elif extracted_coordinates.east.coordinate_value < 2e6 and ground_truth_coordinates["E"] > 2e6:
-                ground_truth_east = int(ground_truth_coordinates["E"]) - 2e6
-                ground_truth_north = int(ground_truth_coordinates["N"]) - 1e6
-            else:
-                ground_truth_east = int(ground_truth_coordinates["E"])
-                ground_truth_north = int(ground_truth_coordinates["N"])
-
-            if (math.isclose(int(extracted_coordinates.east.coordinate_value), ground_truth_east, abs_tol=2)) and (
-                math.isclose(int(extracted_coordinates.north.coordinate_value), ground_truth_north, abs_tol=2)
-            ):
-                coordinate_metrics = Metrics(
-                    tp=1,
-                    fp=0,
-                    fn=0,
-                )
-            else:
-                coordinate_metrics = Metrics(
-                    tp=0,
-                    fp=1,
-                    fn=1,
-                )
+        if extracted_coordinates.east.coordinate_value > 2e6 and ground_truth_coordinates["E"] < 2e6:
+            ground_truth_east = int(ground_truth_coordinates["E"]) + 2e6
+            ground_truth_north = int(ground_truth_coordinates["N"]) + 1e6
+        elif extracted_coordinates.east.coordinate_value < 2e6 and ground_truth_coordinates["E"] > 2e6:
+            ground_truth_east = int(ground_truth_coordinates["E"]) - 2e6
+            ground_truth_north = int(ground_truth_coordinates["N"]) - 1e6
         else:
-            coordinate_metrics = Metrics(
-                tp=0,
-                fp=1 if extracted_coordinates is not None else 0,
-                fn=1 if ground_truth_coordinates is not None else 0,
-            )
+            ground_truth_east = int(ground_truth_coordinates["E"])
+            ground_truth_north = int(ground_truth_coordinates["N"])
 
-        return coordinate_metrics
+        return (math.isclose(int(extracted_coordinates.east.coordinate_value), ground_truth_east, abs_tol=2)) and (
+            math.isclose(int(extracted_coordinates.north.coordinate_value), ground_truth_north, abs_tol=2)
+        )
diff --git a/src/stratigraphy/evaluation/utility.py b/src/stratigraphy/evaluation/utility.py
index a9895c48..4a7212f1 100644
--- a/src/stratigraphy/evaluation/utility.py
+++ b/src/stratigraphy/evaluation/utility.py
@@ -1,32 +1,79 @@
 """Utility functions for evaluation."""
 
-from collections import Counter
+import dataclasses
+from collections.abc import Callable
+from typing import TypeVar
 
 from stratigraphy.evaluation.evaluation_dataclasses import Metrics
 from stratigraphy.layer.layer import LayerDepths
 
 
-def count_against_ground_truth(values: list[any], ground_truth: list[any]) -> Metrics:
-    """Count evaluation metrics by comparing predicted values against ground truth.
+@dataclasses.dataclass
+class EvaluationResults:
+    """Class that contains that data about which extracted entries could be matched with which ground truth values."""
+
+    extracted_correct: list[bool]
+    ground_truth_correct: list[bool]
+
+    @property
+    def metrics(self) -> Metrics:
+        """Derive the metrics (true positives, false positives, fales negatives) from the evaluation results."""
+        tp = sum(1 for is_correct in self.extracted_correct if is_correct)
+        return Metrics(
+            tp=tp,
+            fp=len(self.extracted_correct) - tp,
+            fn=len(self.ground_truth_correct) - tp,
+        )
+
+
+T = TypeVar("T")
+U = TypeVar("U")
+
+
+def evaluate_single(extracted: T | None, ground_truth: U | None, match: Callable[[T, U], bool]) -> EvaluationResults:
+    """Count evaluation metrics by comparing an optional predicted value against an optional ground truth value.
 
-    Metrics are calculated as follows:
-    - True Positives (tp): Number of values that appear in both lists (counting duplicates)
-    - False Positives (fp): Number of extra predictions (len(values) - tp)
-    - False Negatives (fn): Number of missed ground truth values (len(ground_truth) - tp)
+    Args:
+        extracted: (T | None): The predicted value to evaluate (if any).
+        ground_truth (U | None): The ground truth value to compare against (if any).
+        match (Callable[[T, U], bool]): A function that defines when the extracted value matches the ground truth
+            value.
+
+    Returns:
+        EvaluationResults: The evaluation results, including true positive, false positive and false negative counts.
+    """
+    extracted_list = [extracted] if extracted is not None else []
+    ground_truth_list = [ground_truth] if ground_truth is not None else []
+    return evaluate(extracted_list, ground_truth_list, match)
+
+
+def evaluate(extracted: list[T], ground_truth: list[U], match: Callable[[T, U], bool]) -> EvaluationResults:
+    """Count evaluation metrics by comparing predicted values against ground truth.
 
     Args:
-        values (list[any]): The predicted values to evaluate
-        ground_truth (list[any]): The ground truth values to compare against
+        extracted: (list[T]): The predicted values to evaluate.
+        ground_truth (list[U]): The ground truth values to compare against.
+        match (Callable[[T, U], bool]): A function that defines when an extracted value matches a ground truth value.
 
     Returns:
-        Metrics: Object containing tp, fp, and fn counts
+        EvaluationResults: The evaluation results, including true positive, false positive and false negative counts.
     """
-    # Counter deals with duplicates when doing intersection
-    values_counter = Counter(values)
-    ground_truth_counter = Counter(ground_truth)
+    extracted_correct = [False for _ in extracted]
+    ground_truth_correct = [False for _ in ground_truth]
+    for extracted_index, extracted_value in enumerate(extracted):
+        matched_gt_index = next(
+            (
+                ground_truth_index
+                for ground_truth_index, value in enumerate(ground_truth)
+                if not ground_truth_correct[ground_truth_index] and match(extracted_value, value)
+            ),
+            None,
+        )
+        if matched_gt_index is not None:
+            extracted_correct[extracted_index] = True
+            ground_truth_correct[matched_gt_index] = True
 
-    tp = (values_counter & ground_truth_counter).total()  # size of intersection
-    return Metrics(tp=tp, fp=len(values) - tp, fn=len(ground_truth) - tp)
+    return EvaluationResults(extracted_correct, ground_truth_correct)
 
 
 def _is_valid_depth_interval(depths: LayerDepths, start: float, end: float) -> bool:


Test patch:
diff --git a/tests/test_predictions.py b/tests/test_predictions.py
index c4516e7e..28187ddf 100644
--- a/tests/test_predictions.py
+++ b/tests/test_predictions.py
@@ -8,7 +8,7 @@
 from stratigraphy.benchmark.ground_truth import GroundTruth
 from stratigraphy.data_extractor.data_extractor import FeatureOnPage
 from stratigraphy.evaluation.layer_evaluator import LayerEvaluator
-from stratigraphy.evaluation.utility import count_against_ground_truth
+from stratigraphy.evaluation.utility import evaluate, evaluate_single
 from stratigraphy.groundwater.groundwater_extraction import Groundwater, GroundwatersInBorehole
 from stratigraphy.layer.layer import Layer, LayerDepths, LayerDepthsEntry, LayersInBorehole
 from stratigraphy.metadata.coordinate_extraction import CoordinateEntry, LV95Coordinate
@@ -240,7 +240,23 @@ def test_evaluate_metadata_extraction(sample_file_prediction_with_ground_truth:
         ([1, 2], [3, 4], (0, 2, 2)),
     ],
 )
-def test_count_against_ground_truth_cases(values, ground_truth, expected):
+def test_evaluate(values, ground_truth, expected):
     """Test count_against_ground_truth with various scenarios."""
-    metrics = count_against_ground_truth(values, ground_truth)
+    metrics = evaluate(values, ground_truth, lambda a, b: a == b).metrics
+    assert (metrics.tp, metrics.fp, metrics.fn) == expected
+
+
+@pytest.mark.parametrize(
+    "value,ground_truth,expected",
+    [
+        (1, 1, (1, 0, 0)),
+        (1, 2, (0, 1, 1)),
+        (1, None, (0, 1, 0)),
+        (None, 1, (0, 0, 1)),
+        (None, None, (0, 0, 0)),
+    ],
+)
+def test_evaluate_single(value, ground_truth, expected):
+    """Test count_against_ground_truth with various scenarios."""
+    metrics = evaluate_single(value, ground_truth, lambda a, b: a == b).metrics
     assert (metrics.tp, metrics.fp, metrics.fn) == expected