Save all the individual Wasserstein distances

cmd/conf/metrics/default.yaml

@@ -6,6 +6,8 @@ metrics:
     _partial_: true
     random_seed: ${random_seed}
     num_directions: 10000
+    save_all_distances: true
   - _target_: fdiff.sampling.metrics.MarginalWasserstein
     _partial_: true
     random_seed: ${random_seed}
+    save_all_distances: true

src/fdiff/sampling/metrics.py

@@ -1,10 +1,10 @@
 from abc import ABC, abstractmethod, abstractproperty
 from functools import partial
-from pathlib import Path
-from typing import Optional
+from typing import Any, Optional
 
 import numpy as np
 import torch
+
 from fdiff.utils.fourier import dft
 from fdiff.utils.tensors import check_flat_array
 from fdiff.utils.wasserstein import WassersteinDistances
@@ -15,7 +15,7 @@ def __init__(self, original_samples: np.ndarray | torch.Tensor) -> None:
         self.original_samples = check_flat_array(original_samples)
 
     @abstractmethod
-    def __call__(self, other_samples: np.ndarray | torch.Tensor) -> dict[str, float]:
+    def __call__(self, other_samples: np.ndarray | torch.Tensor) -> dict[str, Any]:
         ...
 
     @abstractproperty
@@ -53,7 +53,7 @@ def __init__(
         self.metrics_freq = metrics_freq
         self.include_baselines = include_baselines
 
-    def __call__(self, other_samples: np.ndarray | torch.Tensor) -> dict[str, float]:
+    def __call__(self, other_samples: np.ndarray | torch.Tensor) -> dict[str, Any]:
         metric_dict = {}
         other_samples_freq = dft(other_samples)
         for metric_time, metric_freq in zip(self.metrics_time, self.metrics_freq):
@@ -86,22 +86,27 @@ def __init__(
         original_samples: np.ndarray | torch.Tensor,
         random_seed: int,
         num_directions: int,
+        save_all_distances: bool = False,
     ) -> None:
         super().__init__(original_samples=original_samples)
         self.random_seed = random_seed
         self.num_directions = num_directions
+        self.save_all_distances = save_all_distances
 
-    def __call__(self, other_samples: np.ndarray | torch.Tensor) -> dict[str, float]:
+    def __call__(self, other_samples: np.ndarray | torch.Tensor) -> dict[str, Any]:
         wd = WassersteinDistances(
             original_data=self.original_samples,
             other_data=check_flat_array(other_samples),
             seed=self.random_seed,
         )
         distances = wd.sliced_distances(self.num_directions)
-        return {
+        metrics = {
             "sliced_wasserstein_mean": float(np.mean(distances)),
             "sliced_wasserstein_max": float(np.max(distances)),
         }
+        if self.save_all_distances:
+            metrics["sliced_wasserstein_all"] = distances.tolist()
+        return metrics
 
     @property
     def baseline_metrics(self) -> dict[str, float]:
@@ -141,31 +146,26 @@ def __init__(
         self,
         original_samples: np.ndarray | torch.Tensor,
         random_seed: int,
+        save_all_distances: bool = False,
     ) -> None:
         super().__init__(original_samples=original_samples)
         self.random_seed = random_seed
+        self.save_all_distances = save_all_distances
 
-    def __call__(self, other_samples: np.ndarray | torch.Tensor) -> dict[str, float]:
+    def __call__(self, other_samples: np.ndarray | torch.Tensor) -> dict[str, Any]:
         wd = WassersteinDistances(
             original_data=self.original_samples,
             other_data=check_flat_array(other_samples),
             seed=self.random_seed,
         )
         distances = wd.marginal_distances()
-        return {
+        metrics = {
             "marginal_wasserstein_mean": float(np.mean(distances)),
             "marginal_wasserstein_max": float(np.max(distances)),
         }
-
-    def save(self, other_samples: np.ndarray | torch.Tensor, path: str | Path) -> None:
-        # Save the distances array for post-processing
-        wd = WassersteinDistances(
-            original_data=self.original_samples,
-            other_data=check_flat_array(other_samples),
-            seed=self.random_seed,
-        )
-        distances = wd.marginal_distances()
-        np.save(path, distances)
+        if self.save_all_distances:
+            metrics["marginal_wasserstein_all"] = distances.tolist()
+        return metrics
 
     @property
     def baseline_metrics(self) -> dict[str, float]:

src/fdiff/utils/extraction.py

@@ -91,5 +91,8 @@ def dict_to_str(dict: DictConfig | dict[str, Any]) -> str:
     dict_str = ""
     max_len = max([len(k) for k in dict])
     for k, v in dict.items():
+        # In case of long lists, just print the first 3 elements
+        if isinstance(v, list):
+            v = v[:3] + ["..."] if len(v) > 3 else v
         dict_str += f"\t {k: <{max_len + 5}} : \t  {v} \t \n"
     return dict_str

tests/test_metrics.py

@@ -10,6 +10,7 @@
 max_len = 2
 n_samples = 1000
 num_directions = 1000
+EPS = 1e-5
 
 test_data_wasserstein = [0.0, 0.1, 1.0]
 
@@ -36,9 +37,17 @@ def test_sliced_waserstein(shift: float) -> None:
         original_samples=dataset1,
         random_seed=random_seed,
         num_directions=num_directions,
+        save_all_distances=True,
     )
     metrics = sw(dataset2)
 
+    assert (
+        np.abs(
+            metrics["sliced_wasserstein_mean"]
+            - np.mean(metrics["sliced_wasserstein_all"])
+        )
+        <= EPS
+    )
     assert metrics["sliced_wasserstein_mean"] <= metrics["sliced_wasserstein_max"]
     assert np.abs(metrics["sliced_wasserstein_mean"] - pot_estimate) <= 0.1
 
@@ -56,9 +65,18 @@ def test_marginal_waserstein(shift: float) -> None:
     ground_truth = shift
 
     # Compute sliced wasserstein distance
-    mw = MarginalWasserstein(original_samples=dataset1, random_seed=random_seed)
+    mw = MarginalWasserstein(
+        original_samples=dataset1, random_seed=random_seed, save_all_distances=True
+    )
     metrics = mw(dataset2)
 
+    assert (
+        np.abs(
+            metrics["marginal_wasserstein_mean"]
+            - np.mean(metrics["marginal_wasserstein_all"])
+        )
+        <= EPS
+    )
     assert metrics["marginal_wasserstein_mean"] <= metrics["marginal_wasserstein_max"]
     assert np.abs(metrics["marginal_wasserstein_mean"] - ground_truth) <= 0.1
     assert np.abs(metrics["marginal_wasserstein_max"] - ground_truth) <= 0.1