Updated MMDInterval to use MetricInterval objects

Author: Markus28

polygraph/metrics/base/__init__.py

@@ -3,7 +3,6 @@
     FrechetDistance,
 )
 from polygraph.metrics.base.mmd import (
-    MMDInterval,
     DescriptorMMD2,
     DescriptorMMD2Interval,
     MaxDescriptorMMD2,
@@ -27,7 +26,6 @@
     "MetricInterval",
     "FittedFrechetDistance",
     "FrechetDistance",
-    "MMDInterval",
     "DescriptorMMD2",
     "DescriptorMMD2Interval",
     "MaxDescriptorMMD2",

polygraph/metrics/base/metric_interval.py

@@ -3,6 +3,7 @@
 
 
 class MetricInterval:
+    """Class for representing uncertainty quantifications of a metric."""
     def __init__(
         self,
         mean: float,
@@ -46,6 +47,18 @@ def from_samples(
 
         return cls(mean=mean, std=std, low=low, high=high, coverage=coverage)
 
+    def __getitem__(self, key: str) -> float:
+        if key == "mean":
+            return self.mean
+        elif key == "std":
+            return self.std
+        elif key == "low":
+            return self.low
+        elif key == "high":
+            return self.high
+        else:
+            raise ValueError(f"Invalid key: {key}")
+
     def __repr__(self):
         if self.coverage is not None:
             return f"MetricInterval(mean={self.mean}, std={self.std}, low={self.low}, high={self.high}, coverage={self.coverage})"

polygraph/metrics/base/mmd.py

@@ -48,19 +48,16 @@
 from polygraph.utils.kernels import DescriptorKernel, GramBlocks
 from polygraph.utils.mmd_utils import mmd_from_gram
 from polygraph.metrics.base.interfaces import GenerationMetric, GenerationMetricInterval
+from polygraph.metrics.base.metric_interval import MetricInterval
 
 __all__ = [
     "DescriptorMMD2",
     "MaxDescriptorMMD2",
-    "MMDInterval",
     "DescriptorMMD2Interval",
     "MaxDescriptorMMD2Interval",
 ]
 
 
-MMDInterval = namedtuple("MMDInterval", ["mean", "std", "low", "high"])
-
-
 class DescriptorMMD2(GenerationMetric):
     """Computes squared MMD between reference and generated graphs using a kernel.
 
@@ -201,9 +198,9 @@ def _generate_mmd_samples(
     def compute(
         *args, **kwargs
     ) -> Union[
-        MMDInterval,
+        MetricInterval,
         Dict[str, float],
-        Tuple[Union[MMDInterval, Dict[str, float]], np.ndarray],
+        Tuple[Union[MetricInterval, Dict[str, float]], np.ndarray],
     ]: ...
 
 
@@ -225,13 +222,7 @@ def compute(
         subsample_size: int,
         num_samples: int = 500,
         coverage: float = 0.95,
-        as_scalar_value_dict: bool = False,
-        return_samples: bool = False,
-    ) -> Union[
-        MMDInterval,
-        Dict[str, float],
-        Tuple[Union[MMDInterval, Dict[str, float]], np.ndarray],
-    ]:
+    ) -> MetricInterval:
         """Computes MMD² confidence intervals through subsampling.
 
         Args:
@@ -248,26 +239,8 @@ def compute(
             subsample_size=subsample_size,
             num_samples=num_samples,
         )
-        low, high = (
-            np.quantile(mmd_samples, (1 - coverage) / 2, axis=0),
-            np.quantile(mmd_samples, coverage + (1 - coverage) / 2, axis=0),
-        )
-        avg = np.mean(mmd_samples, axis=0)
-        std = np.std(mmd_samples, axis=0)
-        if as_scalar_value_dict:
-            return_result = {
-                "mean": avg,
-                "std": std,
-                "low": low,
-                "high": high,
-            }
-        else:
-            return_result = MMDInterval(mean=avg, std=std, low=low, high=high)
-
-        if return_samples:
-            return return_result, mmd_samples
-        else:
-            return return_result
+        assert mmd_samples.ndim == 1
+        return MetricInterval.from_samples(mmd_samples, coverage=coverage)
 
 
 class MaxDescriptorMMD2Interval(_DescriptorMMD2Interval, GenerationMetricInterval):
@@ -306,13 +279,7 @@ def compute(
         subsample_size: int,
         num_samples: int = 500,
         coverage: float = 0.95,
-        as_scalar_value_dict: bool = False,
-        return_samples: bool = False,
-    ) -> Union[
-        MMDInterval,
-        Dict[str, float],
-        Tuple[Union[MMDInterval, Dict[str, float]], np.ndarray],
-    ]:
+    ) -> MetricInterval:
         """Computes confidence intervals for maximum MMD² through subsampling.
 
         Args:
@@ -331,24 +298,6 @@ def compute(
         )
         assert mmd_samples.ndim == 2
         mmd_samples = np.max(mmd_samples, axis=1)
-        low, high = (
-            np.quantile(mmd_samples, (1 - coverage) / 2, axis=0),
-            np.quantile(mmd_samples, coverage + (1 - coverage) / 2, axis=0),
-        )
-        avg = np.mean(mmd_samples, axis=0)
-        std = np.std(mmd_samples, axis=0)
-
-        if as_scalar_value_dict:
-            return_result = {
-                "mean": avg,
-                "std": std,
-                "low": low,
-                "high": high,
-            }
-        else:
-            return_result = MMDInterval(mean=avg, std=std, low=low, high=high)
-
-        if return_samples:
-            return return_result, mmd_samples
-        else:
-            return return_result
+        return MetricInterval.from_samples(mmd_samples, coverage=coverage)
+
+

tests/test_mmd.py

@@ -16,7 +16,6 @@
     DescriptorMMD2Interval,
     MaxDescriptorMMD2,
     MaxDescriptorMMD2Interval,
-    MMDInterval,
 )
 from polygraph.metrics.gran import (
     GRANClusteringMMD2,
@@ -49,6 +48,8 @@
 from polygraph.utils.kernels import LinearKernel
 from polygraph.utils.graph_descriptors import WeisfeilerLehmanDescriptor
 from polygraph.utils.mmd_utils import mmd_from_gram
+from polygraph.metrics.base.metric_interval import MetricInterval
+
 import grakel
 
 
@@ -191,7 +192,7 @@ def test_mmd_uncertainty(request, datasets, kernel, subsample_size, variant):
     kernel = request.getfixturevalue(kernel)
     mmd = DescriptorMMD2Interval(sbm, kernel, variant=variant)
     result = mmd.compute(planar, subsample_size=subsample_size)
-    assert isinstance(result, MMDInterval)
+    assert isinstance(result, MetricInterval)
     assert result.std > 0
 
     rng = np.random.default_rng(42)
@@ -243,7 +244,7 @@ def test_concrete_uncertainty(
 
     interval_mmd = interval_cls(planar)
     interval = interval_mmd.compute(sbm, subsample_size=subsample_size)
-    assert isinstance(interval, MMDInterval)
+    assert isinstance(interval, MetricInterval)
 
     num_in_bounds = 0
     num_total = 10
@@ -258,6 +259,7 @@ def test_concrete_uncertainty(
 
         single_mmd = single_cls(planar_samples)
         single_estimate = single_mmd.compute(sbm_samples)
+        assert interval.low <= interval.high
         if interval.low <= single_estimate <= interval.high:
             num_in_bounds += 1