Merge branch 'main' into feature/greedy_kernel_points_analytic_test

Author: tp832944

.pre-commit-config.yaml

@@ -42,7 +42,7 @@ repos:
           - --ignore-case
           - --unique
   - repo: https://github.com/astral-sh/uv-pre-commit
-    rev: 0.6.3
+    rev: 0.6.6
     hooks:
       # Keep lock file up to date
       - id: uv-lock
@@ -105,7 +105,7 @@ repos:
       # Enforce that type annotations are used instead of type comments
       - id: python-use-type-annotations
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.9.9
+    rev: v0.11.0
     hooks:
       # Run the linter.
       - id: ruff

CHANGELOG.md

@@ -11,6 +11,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Added Compress++ coreset reduction algorithm.
   (https://github.com/gchq/coreax/issues/934)
+- Added `reduce_iterative()` method to Kernel Herding. (https://github.com/gchq/coreax/pull/983)
+- Added probabilistic iterative Kernel Herding benchmarking results. (https://github.com/gchq/coreax/pull/983)
 - Analytic example with integration test for `GreedyKernelPoints` plus an analytic unit
   test for the loss function. (https://github.com/gchq/coreax/pull/1004)
 
@@ -20,6 +22,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- Changed the score function used by Stein Thinning in benchmarking.
+  (https://github.com/gchq/coreax/pull/1000)
+- Fixed the random state for UMAP in benchmarking for reproducibility.
+  (https://github.com/gchq/coreax/pull/1000)
+- Reduced the number of dimensions when applying UMAP in `pounce_benchmark.py`.
+  (https://github.com/gchq/coreax/pull/1000)
 - Refactored `GreedyKernelPoints` and associated functions to make more extensible in
   future. (https://github.com/gchq/coreax/pull/1004)
 

benchmark/blobs_benchmark.py

@@ -31,13 +31,17 @@
 import json
 import os
 import time
+from typing import Union
 
 import jax
 import jax.numpy as jnp
+import jax.scipy as jsp
 import numpy as np
+from jaxtyping import Array, Shaped
 from sklearn.datasets import make_blobs
 
-from coreax import Data, SlicedScoreMatching
+from coreax import Data
+from coreax.benchmark_util import IterativeKernelHerding
 from coreax.kernels import (
     SquaredExponentialKernel,
     SteinKernel,
@@ -46,7 +50,6 @@
 from coreax.metrics import KSD, MMD
 from coreax.solvers import (
     CompressPlusPlus,
-    IterativeKernelHerding,
     KernelHerding,
     KernelThinning,
     RandomSample,
@@ -84,17 +87,32 @@ def setup_stein_kernel(
     :param random_seed: An integer seed for the random number generator.
     :return: A SteinKernel object.
     """
-    sliced_score_matcher = SlicedScoreMatching(
-        jax.random.PRNGKey(random_seed),
-        jax.random.rademacher,
-        use_analytic=True,
-        num_random_vectors=100,
-        learning_rate=0.001,
-        num_epochs=50,
-    )
+    # Fit a Gaussian kernel density estimator on a subset of points for efficiency
+    num_data_points = len(dataset)
+    num_samples_length_scale = min(num_data_points, 1000)
+    generator = np.random.default_rng(random_seed)
+    idx = generator.choice(num_data_points, num_samples_length_scale, replace=False)
+    kde = jsp.stats.gaussian_kde(dataset.data[idx].T)
+
+    # Define the score function as the gradient of log density given by the KDE
+    def score_function(
+        x: Union[Shaped[Array, " n d"], Shaped[Array, ""], float, int],
+    ) -> Union[Shaped[Array, " n d"], Shaped[Array, " 1 1"]]:
+        """
+        Compute the score function (gradient of log density) for a single point.
+
+        :param x: Input point represented as array
+        :return: Gradient of log probability density at the given point
+        """
+
+        def logpdf_single(x: Shaped[Array, " d"]) -> Shaped[Array, ""]:
+            return kde.logpdf(x.reshape(1, -1))[0]
+
+        return jax.grad(logpdf_single)(x)
+
     return SteinKernel(
         base_kernel=sq_exp_kernel,
-        score_function=sliced_score_matcher.match(jnp.asarray(dataset.data)),
+        score_function=score_function,
     )
 
 
@@ -142,7 +160,7 @@ def setup_solvers(
             SteinThinning(
                 coreset_size=coreset_size,
                 kernel=stein_kernel,
-                regularise=False,
+                regularise=True,
             ),
         ),
         (
@@ -188,6 +206,18 @@ def setup_solvers(
                 num_iterations=5,
             ),
         ),
+        (
+            "CubicProbIterativeHerding",
+            IterativeKernelHerding(
+                coreset_size=coreset_size,
+                kernel=sq_exp_kernel,
+                probabilistic=True,
+                temperature=0.001,
+                random_key=random_key,
+                num_iterations=10,
+                t_schedule=1 / jnp.linspace(10, 100, 10) ** 3,
+            ),
+        ),
     ]
 
 
@@ -296,7 +326,7 @@ def main() -> None:  # pylint: disable=too-many-locals
 
         # Set up metrics
         mmd_metric = MMD(kernel=sq_exp_kernel)
-        ksd_metric = KSD(kernel=sq_exp_kernel)
+        ksd_metric = KSD(kernel=stein_kernel)  # KSD needs a Stein kernel
 
         # Set up weights optimiser
         weights_optimiser = MMDWeightsOptimiser(kernel=sq_exp_kernel)