Refactor recall calculation and add type hints

ham.py

@@ -1,18 +1,18 @@
 """This module provides functions for calculating hierarchical variants of precicion, recall and F1."""
 
-from typing import List, Set, Dict, Tuple
+from typing import List, Dict, Tuple, Set
 
 
-def find_ancestors(node: str, hierarchy: dict) -> set:
+def find_ancestors(node: str, hierarchy: Dict[str, Set[str]]) -> Set[str]:
     """
     Find the ancestors of a given node in a hierarchy.
 
     Args:
         node (str): The node for which to find ancestors.
-        hierarchy (dict): A dictionary representing the hierarchy, where the keys are nodes and the values are their parents.
+        hierarchy (Dict[str, Set[str]]): A dictionary representing the hierarchy, where the keys are nodes and the values are their parents.
 
     Returns:
-        set: A set of ancestors of the given node.
+        Set[str]: A set of ancestors of the given node.
     """
     ancestors = set()
     nodes_to_visit = [node]
@@ -54,45 +54,52 @@ def calculate_hierarchical_precision_recall(
     Args:
         reference_codes (List[str]): The list of reference codes.
         predicted_codes (List[str]): The list of predicted codes.
-        hierarchy (Dict[str, Set[str]]): The hierarchy definition where keys are nodes and values are sets of parent nodes.
+        hierarchy (Dict[str, Dict[str, float]]): The hierarchy definition where keys are nodes and values are dictionaries of parent nodes with distances.
 
     Returns:
         Tuple[float, float]: A tuple containing the hierarchical precision and recall floating point values.
     """
     extended_real = {}
+    extended_predicted = {}
 
     # Extend the sets of reference codes with their ancestors
     for code in reference_codes:
-        weight = 1.0  # Full weight for exact match
-        extended_real[code] = weight
+        extended_real[code] = 1.0  # Full weight for exact match
         for ancestor, ancestor_weight in hierarchy.get(code, {}).items():
             extended_real[ancestor] = max(
                 extended_real.get(ancestor, 0), ancestor_weight
             )
 
-    extended_predicted = {}
-
     # Extend the sets of predicted codes with their ancestors
     for code in predicted_codes:
-        weight = 1.0
-        extended_predicted[code] = weight
+        extended_predicted[code] = 1.0
         for ancestor, ancestor_weight in hierarchy.get(code, {}).items():
             extended_predicted[ancestor] = max(
                 extended_predicted.get(ancestor, 0), ancestor_weight
             )
 
-    # Calculate weighted correct predictions
-    correct_weights = 0
+    # Calculate weighted correct predictions for precision
+    correct_weights_precision = 0
     for code, weight in extended_predicted.items():
         if code in extended_real:
-            correct_weights += min(weight, extended_real[code])
+            correct_weights_precision += min(weight, extended_real[code])
+
+    # Calculate weighted correct predictions for recall
+    correct_weights_recall = 0
+    for code, weight in extended_real.items():
+        if code in extended_predicted:
+            correct_weights_recall += min(weight, extended_predicted[code])
 
     total_predicted_weights = sum(extended_predicted.values())
     total_real_weights = sum(extended_real.values())
 
     # Calculate hierarchical precision and recall using weighted sums
-    hP = correct_weights / total_predicted_weights if total_predicted_weights else 0
-    hR = correct_weights / total_real_weights if total_real_weights else 0
+    hP = (
+        correct_weights_precision / total_predicted_weights
+        if total_predicted_weights
+        else 0
+    )
+    hR = correct_weights_recall / total_real_weights if total_real_weights else 0
 
     return hP, hR