refactor to arranged code

matching_series.py

@@ -164,11 +164,14 @@ class matching_series(evaluate.Metric):
             return_coverages = True
         predictions = np.array(predictions).astype(dtype)
         references = np.array(references).astype(dtype)
+
         if instance_normalization:
             predictions = (predictions - predictions.mean(axis=1, keepdims=True)) / predictions.std(
                 axis=1, keepdims=True
             )
             references = (references - references.mean(axis=1, keepdims=True)) / references.std(axis=1, keepdims=True)
+
+        assert isinstance(predictions, np.ndarray) and isinstance(references, np.ndarray)
         if predictions.shape[1:] != references.shape[1:]:
             raise ValueError(
                 "The number of features in the predictions and references should be the same. predictions: {}, references: {}".format(
@@ -177,7 +180,110 @@ class matching_series(evaluate.Metric):
             )
 
         # at first, convert the inputs to numpy arrays
+        distance = self.compute_distance(
+            predictions=predictions,
+            references=references,
+            metric=metric,
+            batch_size=batch_size,
+            num_process=num_process,
+            dtype=dtype,
+        )
+
+        metrics = self._compute_metrics(
+            distance=distance.mean(axis=-1),
+            eps=eps,
+            cuc_n_calculation=cuc_n_calculation,
+            cuc_n_samples=cuc_n_samples,
+        )
+        metrics_feature = [
+            self._compute_metrics(distance[:, :, f], eps, cuc_n_calculation, cuc_n_samples)
+            for f in range(predictions.shape[-1])
+        ]
+        macro_metrics = {
+            "macro_" + k: statistics.mean([m[k] for m in metrics_feature])  # type: ignore
+            for k in metrics_feature[0].keys()
+            if isinstance(metrics_feature[0][k], (int, float))
+        }
+
+        out = {}
+        out.update({k: v for k, v in metrics.items() if isinstance(v, (int, float))})
+        out.update(macro_metrics)
+
+        if return_distance:
+            out["distance"] = distance
+        if return_matching:
+            out.update({k: v for k, v in metrics.items() if "match" in k})
+        if return_coverages:
+            out["coverages"] = metrics["coverages"]
+        if return_each_features:
+            out.update(
+                {
+                    k + "_features": [m[k] for m in metrics_feature]
+                    for k in metrics_feature[0].keys()
+                    if isinstance(metrics_feature[0][k], (int, float))
+                }
+            )
+            if return_coverages:
+                out.update(
+                    {
+                        "coverages_features": [m["coverages"] for m in metrics_feature],
+                    }
+                )
+        return out
+
+    def compute_cuc(
+        self,
+        match: np.ndarray,
+        n_reference: int,
+        n_calculation: int,
+        n_samples: Union[List[int], str],
+    ):
+        """
+        Compute Coverage Under Curve
+        Args:
+            match: best match for each generated time series
+            n_reference: number of reference time series
+            n_calculation: number of Coverage Under Curve calculate times
+            n_samples: number of samples to use for Coverage Under Curve calculation. If "auto", it uses the number of samples of the predictions.
+        Returns:
+        """
+        n_generaiton = len(match)
+        if n_samples == "auto":
+            exp = int(math.log2(n_generaiton))
+            n_samples = [int(2**i) for i in range(exp)]
+            n_samples.append(n_generaiton)
+        assert isinstance(n_samples, list) and all(isinstance(n, int) for n in n_samples)
 
+        coverages = []
+        for n_sample in n_samples:
+            coverage = 0
+            for _ in range(n_calculation):
+                sample = np.random.choice(match, size=n_sample, replace=False)  # type: ignore
+                coverage += len(np.unique(sample)) / n_reference
+            coverages.append(coverage / n_calculation)
+        cuc = (np.trapz(coverages, n_samples) / len(n_samples) / max(n_samples)).item()
+        return coverages, cuc
+
+    @staticmethod
+    def _compute_distance(x, y, metric: str = "mse", axis: int = -1):
+        if metric.lower() == "mse":
+            return np.mean((x - y) ** 2, axis=axis)
+        elif metric.lower() == "mae":
+            return np.mean(np.abs(x - y), axis=axis)
+        elif metric.lower() == "rmse":
+            return np.sqrt(np.mean((x - y) ** 2, axis=axis))
+        else:
+            raise ValueError("Unknown metric: {}".format(metric))
+
+    def compute_distance(
+        self,
+        predictions: np.ndarray,
+        references: np.ndarray,
+        metric: str,
+        batch_size: Optional[int] = None,
+        num_process: int = 1,
+        dtype=np.float32,
+    ):
         # distance between predictions and references for all example combinations for each features
         # shape: (num_generation, num_reference, num_features)
         if batch_size is not None:
@@ -194,7 +300,7 @@ class matching_series(evaluate.Metric):
                 ]
                 with concurrent.futures.ProcessPoolExecutor(max_workers=num_process) as executor:
                     results = executor.map(
-                        self._compute_metric,
+                        self._compute_distance,
                         *zip(*args),
                     )
                     for (i, j), d in zip(idxs, results):
@@ -205,7 +311,7 @@ class matching_series(evaluate.Metric):
                 # iterate over the predictions and references in batches
                 for i in range(0, len(predictions) + batch_size, batch_size):
                     for j in range(0, len(references) + batch_size, batch_size):
-                        d = self._compute_metric(
+                        d = self._compute_distance(
                             predictions[i : i + batch_size, None],
                             references[None, j : j + batch_size],
                             metric=metric,
@@ -213,24 +319,34 @@ class matching_series(evaluate.Metric):
                         )
                         distance[i : i + batch_size, j : j + batch_size] = d
         else:
-            distance = self._compute_metric(predictions[:, None], references[None, :], metric=metric, axis=-2)
+            distance = self._compute_distance(predictions[:, None], references[None, :], metric=metric, axis=-2)
+        return distance
 
+    def _compute_metrics(
+        self,
+        distance: np.ndarray,
+        eps: float = 1e-10,
+        cuc_n_calculation: int = 3,
+        cuc_n_samples: Union[List[int], str] = "auto",
+    ) -> dict[str, float | list[float]]:
+        """
+        Compute metrics from the distance matrix
+        Args:
+            distance: distance matrix. shape: (num_generation, num_reference)
+        Returns:
+        """
         index_distance = distance.diagonal(axis1=0, axis2=1).mean().item()
 
         # matching scores
-        distance_mean = distance.mean(axis=-1)
         # best match for each generated time series
         # shape: (num_generation,)
-        best_match = np.argmin(distance_mean, axis=-1)
-
-        # matching distance
-        # shape: (num_generation,)
-        precision_distance = distance_mean[np.arange(len(best_match)), best_match].mean().item()
+        best_match = np.argmin(distance, axis=-1)
+        precision_distance = distance[np.arange(len(best_match)), best_match].mean().item()
 
         # best match for each reference time series
         # shape: (num_reference,)
-        best_match_inv = np.argmin(distance_mean, axis=0)
-        recall_distance = distance_mean[best_match_inv, np.arange(len(best_match_inv))].mean().item()
+        best_match_inv = np.argmin(distance, axis=0)
+        recall_distance = distance[best_match_inv, np.arange(len(best_match_inv))].mean().item()
 
         f1_distance = 2 / (1 / (precision_distance + eps) + 1 / (recall_distance + eps))
         mean_distance = (precision_distance + recall_distance) / 2
@@ -240,144 +356,19 @@ class matching_series(evaluate.Metric):
         matching_precision = np.unique(best_match_inv).size / len(best_match)
         matching_f1 = 2 / (1 / (matching_precision + eps) + 1 / (matching_recall + eps))
 
-        # take matching for each feature and compute metrics for them
-        precision_distance_features = []
-        recall_distance_features = []
-        f1_distance_features = []
-        mean_distance_features = []
-        matching_precision_features = []
-        matching_recall_features = []
-        matching_f1_features = []
-        index_distance_features = []
-        coverages_features = []
-        cuc_features = []
-        for f in range(predictions.shape[-1]):
-            distance_f = distance[:, :, f]
-            index_distance_f = (distance_f.diagonal(axis1=0, axis2=1).mean()).item()
-            best_match_f = np.argmin(distance_f, axis=-1)
-            precision_distance_f = (distance_f[np.arange(len(best_match_f)), best_match_f].mean()).item()
-            best_match_inv_f = np.argmin(distance_f, axis=0)
-            recall_distance_f = (distance_f[best_match_inv_f, np.arange(len(best_match_inv_f))].mean()).item()
-            f1_distance_f = 2 / (1 / (precision_distance_f + eps) + 1 / (recall_distance_f + eps))
-            mean_distance_f = (precision_distance_f + recall_distance_f) / 2
-            precision_distance_features.append(precision_distance_f)
-            recall_distance_features.append(recall_distance_f)
-            f1_distance_features.append(f1_distance_f)
-            index_distance_features.append(index_distance_f)
-            mean_distance_features.append(mean_distance_f)
-
-            matching_recall_f = np.unique(best_match_f).size / len(best_match_f)
-            matching_precision_f = np.unique(best_match_inv_f).size / len(best_match_inv_f)
-            matching_f1_f = 2 / (1 / (matching_precision_f + eps) + 1 / (matching_recall_f + eps))
-            matching_precision_features.append(matching_precision_f)
-            matching_recall_features.append(matching_recall_f)
-            matching_f1_features.append(matching_f1_f)
-
-            coverages_f, cuc_f = self.compute_cuc(best_match_f, len(references), cuc_n_calculation, cuc_n_samples)
-            coverages_features.append(coverages_f)
-            cuc_features.append(cuc_f)
-
-        macro_precision_distance = statistics.mean(precision_distance_features)
-        macro_recall_distance = statistics.mean(recall_distance_features)
-        macro_f1_distance = statistics.mean(f1_distance_features)
-        macro_mean_distance = statistics.mean(mean_distance_features)
-        macro_index_distance = statistics.mean(index_distance_features)
-
-        macro_matching_precision = statistics.mean(matching_precision_features)
-        macro_matching_recall = statistics.mean(matching_recall_features)
-        macro_matching_f1 = statistics.mean(matching_f1_features)
-
         # cuc
-        coverages, cuc = self.compute_cuc(best_match, len(references), cuc_n_calculation, cuc_n_samples)
-
-        macro_cuc = statistics.mean(cuc_features)
-        macro_coverages = [statistics.mean(c) for c in zip(*coverages_features)]
-        out = {
+        coverages, cuc = self.compute_cuc(best_match, len(best_match_inv), cuc_n_calculation, cuc_n_samples)
+        return {
             "precision_distance": precision_distance,
             "f1_distance": f1_distance,
             "recall_distance": recall_distance,
             "mean_distance": mean_distance,
             "index_distance": index_distance,
-            "macro_precision_distance": macro_precision_distance,
-            "macro_recall_distance": macro_recall_distance,
-            "macro_f1_distance": macro_f1_distance,
-            "macro_mean_distance": macro_mean_distance,
-            "macro_index_distance": macro_index_distance,
             "matching_precision": matching_precision,
             "matching_recall": matching_recall,
             "matching_f1": matching_f1,
-            "macro_matching_precision": macro_matching_precision,
-            "macro_matching_recall": macro_matching_recall,
-            "macro_matching_f1": macro_matching_f1,
             "cuc": cuc,
-            "macro_cuc": macro_cuc,
+            "coverages": coverages,
+            "match": best_match,
+            "match_inv": best_match_inv,
         }
-        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,
-        match: np.ndarray,
-        n_reference: int,
-        n_calculation: int,
-        n_samples: Union[List[int], str],
-    ):
-        """
-        Compute Coverage Under Curve
-        Args:
-            match: best match for each generated time series
-            n_reference: number of reference time series
-            n_calculation: number of Coverage Under Curve calculate times
-            n_samples: number of samples to use for Coverage Under Curve calculation. If "auto", it uses the number of samples of the predictions.
-        Returns:
-        """
-        n_generaiton = len(match)
-        if n_samples == "auto":
-            exp = int(math.log2(n_generaiton))
-            n_samples = [int(2**i) for i in range(exp)]
-            n_samples.append(n_generaiton)
-        assert isinstance(n_samples, list) and all(isinstance(n, int) for n in n_samples)
-
-        coverages = []
-        for n_sample in n_samples:
-            coverage = 0
-            for _ in range(n_calculation):
-                sample = np.random.choice(match, size=n_sample, replace=False)  # type: ignore
-                coverage += len(np.unique(sample)) / n_reference
-            coverages.append(coverage / n_calculation)
-        cuc = (np.trapz(coverages, n_samples) / len(n_samples) / max(n_samples)).item()
-        return coverages, cuc
-
-    @staticmethod
-    def _compute_metric(x, y, metric: str = "mse", axis: int = -1):
-        if metric.lower() == "mse":
-            return np.mean((x - y) ** 2, axis=axis)
-        elif metric.lower() == "mae":
-            return np.mean(np.abs(x - y), axis=axis)
-        elif metric.lower() == "rmse":
-            return np.sqrt(np.mean((x - y) ** 2, axis=axis))
-        else:
-            raise ValueError("Unknown metric: {}".format(metric))