make intermediate results optional for report

matching_series.py

@@ -136,6 +136,11 @@ class matching_series(evaluate.Metric):
         cuc_n_samples: Union[List[int], str] = "auto",
         metric: str = "mse",
         num_process: int = 1,
+        return_distance: bool = False,
+        return_matching: bool = False,
+        return_each_features: bool = False,
+        return_coverages: bool = False,
+        return_all: bool = False,
     ):
         """
         Compute the scores of the module given the predictions and references
@@ -149,6 +154,11 @@ class matching_series(evaluate.Metric):
             cuc_n_samples: number of samples to use for Coverage Under Curve calculation. If "auto", it uses the number of samples of the predictions.
         Returns:
         """
+        if return_all:
+            return_distance = True
+            return_matching = True
+            return_each_features = True
+            return_coverages = True
         predictions = np.array(predictions)
         references = np.array(references)
         if predictions.shape[1:] != references.shape[1:]:
@@ -271,16 +281,11 @@ class matching_series(evaluate.Metric):
 
         macro_cuc = statistics.mean(cuc_features)
         macro_coverages = [statistics.mean(c) for c in zip(*coverages_features)]
-
-        return {
+        out = {
             "precision_distance": precision_distance,
             "f1_distance": f1_distance,
             "recall_distance": recall_distance,
             "index_distance": index_distance,
-            "precision_distance_features": precision_distance_features,
-            "f1_distance_features": f1_distance_features,
-            "recall_distance_features": recall_distance_features,
-            "index_distance_features": index_distance_features,
             "macro_precision_distance": macro_precision_distance,
             "macro_recall_distance": macro_recall_distance,
             "macro_f1_distance": macro_f1_distance,
@@ -288,19 +293,37 @@ class matching_series(evaluate.Metric):
             "matching_precision": matching_precision,
             "matching_recall": matching_recall,
             "matching_f1": matching_f1,
-            "matching_precision_features": matching_precision_features,
-            "matching_recall_features": matching_recall_features,
-            "matching_f1_features": matching_f1_features,
             "macro_matching_precision": macro_matching_precision,
             "macro_matching_recall": macro_matching_recall,
             "macro_matching_f1": macro_matching_f1,
             "cuc": cuc,
-            "coverages": coverages,
             "macro_cuc": macro_cuc,
-            "macro_coverages": macro_coverages,
-            "cuc_features": cuc_features,
-            "coverages_features": coverages_features,
         }
+        if return_distance:
+            out["distance"] = distance
+        if return_matching:
+            out["match"] = best_match
+            out["match_inv"] = best_match_inv
+        if return_each_features:
+            if return_distance:
+                out["distance_features"] = distance_mean
+            out.update(
+                {
+                    "precision_distance_features": precision_distance_features,
+                    "f1_distance_features": f1_distance_features,
+                    "recall_distance_features": recall_distance_features,
+                    "index_distance_features": index_distance_features,
+                    "matching_precision_features": matching_precision_features,
+                    "matching_recall_features": matching_recall_features,
+                    "matching_f1_features": matching_f1_features,
+                    "cuc_features": cuc_features,
+                    "coverages_features": coverages_features,
+                }
+            )
+        if return_coverages:
+            out["coverages"] = coverages
+            out["macro_coverages"] = macro_coverages
+        return out
 
     def compute_cuc(
         self,