Compute kernels in HammingGraph, HypercubeGraph in the log domain

.github/workflows/code-style.yaml

@@ -22,7 +22,7 @@ jobs:
     - name: Install uv
       uses: astral-sh/setup-uv@v6
       with:
-        version: "0.9.6"
+        version: "0.11.23"
       #----------------------------------------------
       #       install
       #----------------------------------------------

.github/workflows/docs.yaml

@@ -16,7 +16,7 @@ jobs:
     - uses: actions/checkout@v4
     - uses: astral-sh/setup-uv@v6
       with:
-        version: "0.9.6"
+        version: "0.11.23"
       #----------------------------------------------
       #       install
       #----------------------------------------------

.github/workflows/quality-checks.yaml

@@ -16,7 +16,7 @@ jobs:
     strategy:
       matrix:
         os: [ubuntu-latest, macos-latest, windows-latest]
-        python-version: ["3.10", "3.11", "3.12"]
+        python-version: ["3.11", "3.12", "3.13"]
       fail-fast: false
 
     runs-on: ${{ matrix.os }}
@@ -30,7 +30,7 @@ jobs:
     - name: Install uv
       uses: astral-sh/setup-uv@v6
       with:
-        version: "0.9.6"
+        version: "0.11.23"
     - name: Install make
       if: runner.os == 'Windows'
       run: choco install make -y

Makefile

@@ -43,6 +43,5 @@ lint: sync
 
 test: sync ## Run the tests, start with the failing ones and break on first fail.
 	@$(UV_RUN) pytest -v -x --ff -rN -Wignore -s --tb=short --durations=0 --cov --cov-report=xml --cov-report=html:coverage_html tests
-	# gpflow is ignored due to incompatibility with the recent setuptools
-	@$(UV_RUN) pytest --nbmake --nbmake-kernel=python3 --durations=0 --nbmake-timeout=1000 --ignore=notebooks/frontends/GPflow.ipynb notebooks/
+	@$(UV_RUN) pytest --nbmake --nbmake-kernel=python3 --durations=0 --nbmake-timeout=1000 notebooks/
 	@echo -e "$(SUCCESS)Tests done$(RESET)"

docs/index.rst

@@ -203,7 +203,7 @@ To install JAX, follow `these instructions <https://github.com/google/jax#instal
 
 ..  code-block:: bash
 
-   pip install gpjax
+   pip install "gpjax>=0.14.0"
 
 .. raw:: html
 

geometric_kernels/feature_maps/__init__.py

@@ -15,6 +15,9 @@
     RandomPhaseFeatureMapCompact,
     RandomPhaseFeatureMapNoncompact,
 )
+from geometric_kernels.feature_maps.random_phase_log_domain import (
+    RandomPhaseFeatureMapLogDomain,
+)
 from geometric_kernels.feature_maps.rejection_sampling import (
     RejectionSamplingFeatureMapHyperbolic,
     RejectionSamplingFeatureMapSPD,

geometric_kernels/feature_maps/random_phase_log_domain.py

@@ -0,0 +1,83 @@
+"""Log-domain random-phase features for discrete-spectrum spaces."""
+
+import lab as B
+from beartype.typing import Dict, Tuple
+
+from geometric_kernels.feature_maps.random_phase import RandomPhaseFeatureMapCompact
+from geometric_kernels.lab_extras import from_numpy, is_complex
+from geometric_kernels.spaces import DiscreteSpectrumSpace
+
+
+class RandomPhaseFeatureMapLogDomain(RandomPhaseFeatureMapCompact):
+    """Random-phase features with log-domain spectral weighting.
+
+    Sampling, feature ordering, and row normalization follow
+    :class:`RandomPhaseFeatureMapCompact`. When the eigenfunctions support log
+    products, the map avoids premature spectral underflow and large
+    multiplicities. Otherwise it uses the standard compact feature map.
+
+    :param space:
+        A discrete-spectrum space with random sampling and addition-theorem
+        eigenfunctions.
+    :param num_levels:
+        Number of spectral levels to include.
+    :param num_random_phases:
+        Number of sampled phases. The map returns this many features per level.
+    """
+
+    def __init__(
+        self,
+        space: DiscreteSpectrumSpace,
+        num_levels: int,
+        num_random_phases: int = 3000,
+    ):
+        super().__init__(space, num_levels, num_random_phases)
+
+    def __call__(
+        self,
+        X: B.Numeric,
+        params: Dict[str, B.Numeric],
+        *,
+        key: B.RandomState,
+        normalize: bool = True,
+        **kwargs,
+    ) -> Tuple[B.RandomState, B.Numeric]:
+        """Return the updated random key and log-domain random-phase features.
+
+        Arguments and return shapes follow ``RandomPhaseFeatureMapCompact``.
+        Normalization produces unit-norm feature rows; unnormalized features
+        can still exceed the floating-point range.
+        """
+        if not self.eigenfunctions.supports_log_domain:
+            return super().__call__(X, params, key=key, normalize=normalize, **kwargs)
+
+        from geometric_kernels.kernels.karhunen_loeve_log_domain import (
+            MaternKarhunenLoeveLogDomain,
+        )
+
+        key, phases = self.space.random(key, self.num_random_phases)
+        log_spectrum = MaternKarhunenLoeveLogDomain.log_spectrum(
+            self.space.get_eigenvalues(self.num_levels),
+            params["nu"],
+            params["lengthscale"],
+            self.space.dimension,
+        )
+        phases = B.cast(B.dtype(X), from_numpy(X, phases))
+        return key, self._features_from_log_spectrum(log_spectrum, X, phases, normalize)
+
+    def _features_from_log_spectrum(self, log_spectrum, X, phases, normalize):
+        log_phi_magnitude, signs = self.eigenfunctions.phi_product_log(
+            X, phases, dtype=B.dtype(log_spectrum)
+        )
+        log_magnitude = log_phi_magnitude + B.transpose(0.5 * log_spectrum)
+        log_magnitude = B.reshape(log_magnitude, X.shape[0], -1)
+        signs = B.reshape(signs, X.shape[0], -1)
+        if normalize:
+            log_magnitude = log_magnitude - B.max(log_magnitude, axis=1, squeeze=False)
+            log_magnitude = log_magnitude - 0.5 * B.logsumexp(
+                2 * log_magnitude, axis=1, squeeze=False
+            )
+        features = signs * B.exp(log_magnitude)
+        if is_complex(features):
+            features = B.concat(B.real(features), B.imag(features), axis=1)
+        return features

geometric_kernels/frontends/gpjax.py

@@ -6,14 +6,12 @@
 :doc:`frontends/GPJax.ipynb </examples/frontends/GPJax>` notebook.
 """
 
-from dataclasses import dataclass
-
+import equinox as eqx
 import gpjax
 import jax.numpy as jnp
+import lineax as lx
+import paramax
 from beartype.typing import List, TypeVar, Union
-from flax import nnx
-from gpjax.kernels.computations.base import AbstractKernelComputation
-from gpjax.linalg import Diagonal, psd
 from gpjax.parameters import NonNegativeReal, PositiveReal
 from gpjax.typing import Array, ScalarFloat
 from jaxtyping import Float, Num
@@ -49,7 +47,8 @@ def cross_covariance(
         :return:
             The N x M covariance matrix.
         """
-        nu_value = kernel.nu.value if kernel.trainable_nu else kernel.nu
+        kernel = paramax.unwrap(kernel)
+        nu_value = kernel.nu
 
         # Ensure inputs have `ndim` > 1. GPJax may squeeze shape `(1, 1)` into
         # `(1,)` which causes issues when passing to the base kernel.
@@ -58,13 +57,13 @@ def cross_covariance(
         if y.ndim == 1:
             y = y[:, jnp.newaxis]
 
-        return kernel.variance.value * kernel.base_kernel.K(
-            {"lengthscale": kernel.lengthscale.value, "nu": nu_value}, x, y
+        return kernel.variance * kernel.base_kernel.K(
+            {"lengthscale": kernel.lengthscale, "nu": nu_value}, x, y
         )
 
     def diagonal(
         self, kernel: Kernel, x: Num[Array, "N #D1 D2"]  # noqa: F821
-    ) -> Diagonal:
+    ) -> lx.AbstractLinearOperator:
         """
         Compute the diagonal of the covariance matrix `K(x, x)` where `x` is a batch of
         vectors (or a batch of matrices) of inputs.
@@ -79,24 +78,25 @@ def diagonal(
         Returns:
             The computed diagonal variance as a `Diagonal` linear operator.
         """
-        nu_value = kernel.nu.value if kernel.trainable_nu else kernel.nu
+        kernel = paramax.unwrap(kernel)
+        nu_value = kernel.nu
 
         # Ensure inputs have `ndim` > 1. GPJax may squeeze shape `(1, 1)` into
         # `(1,)` which causes issues when passing to the base kernel.
         if x.ndim == 1:
             x = x[:, jnp.newaxis]
 
-        return psd(
-            Diagonal(
-                kernel.variance.value
+        return lx.TaggedLinearOperator(
+            lx.DiagonalLinearOperator(
+                kernel.variance
                 * kernel.base_kernel.K_diag(
-                    {"lengthscale": kernel.lengthscale.value, "nu": nu_value}, x
+                    {"lengthscale": kernel.lengthscale, "nu": nu_value}, x
                 )
-            )
+            ),
+            lx.positive_semidefinite_tag,
         )
 
 
-@dataclass
 class GPJaxGeometricKernel(gpjax.kernels.AbstractKernel):
     r"""
     GPJax wrapper for :class:`~.kernels.BaseGeometricKernel`.
@@ -108,9 +108,8 @@ class GPJaxGeometricKernel(gpjax.kernels.AbstractKernel):
     .. note::
         Remember that the `base_kernel` itself does not store any of its
         hyperparameters (like `lengthscale` and `nu`). If you do not set them
-        manually—when initializing the object or after, by setting the
-        properties—this wrapper will use the values provided by
-        `base_kernel.init_params`.
+        manually when initializing the object, this wrapper will use the values
+        provided by `base_kernel.init_params`.
 
     :param base_kernel:
         The kernel to wrap.
@@ -135,29 +134,31 @@ class GPJaxGeometricKernel(gpjax.kernels.AbstractKernel):
         Defaults to False.
     """
 
-    nu: Union[ScalarFloat, nnx.Variable[ScalarFloat], None]
-    lengthscale: nnx.Variable[Union[ScalarFloat, Float[Array, " D"]]]
-    variance: nnx.Variable[ScalarFloat]
+    nu: paramax.AbstractUnwrappable
+    lengthscale: paramax.AbstractUnwrappable
+    variance: paramax.AbstractUnwrappable
 
-    base_kernel: BaseGeometricKernel
-    compute_engine: AbstractKernelComputation = _GeometricKernelComputation()
-    name: str = "Geometric Kernel"
+    base_kernel: BaseGeometricKernel = eqx.field(static=True)
+    trainable_nu: bool = eqx.field(static=True)
+    name: str = eqx.field(static=True, default="Geometric Kernel")
 
     def __init__(
         self,
         base_kernel: BaseGeometricKernel,
         lengthscale: Union[
             Union[ScalarFloat, Float[Array, " D"]],
-            nnx.Variable[Union[ScalarFloat, Float[Array, " D"]]],
+            paramax.AbstractUnwrappable,
             None,
         ] = None,
-        nu: Union[ScalarFloat, nnx.Variable[ScalarFloat], None] = None,
-        variance: Union[ScalarFloat, nnx.Variable[ScalarFloat]] = 1.0,
+        nu: Union[ScalarFloat, paramax.AbstractUnwrappable, None] = None,
+        variance: Union[ScalarFloat, paramax.AbstractUnwrappable] = 1.0,
         trainable_nu: bool = False,
     ):
-        active_dims = None
-        n_dims = None
-        super().__init__(active_dims, n_dims, self.compute_engine)
+        # Initialise inherited fields directly: Equinox freezes the module when
+        # a parent constructor returns.
+        self.active_dims = slice(None)
+        self.n_dims = None
+        self.compute_engine = _GeometricKernelComputation()
 
         self.base_kernel = base_kernel
         default_params = self.base_kernel.init_params()
@@ -167,20 +168,20 @@ def __init__(
         if nu is None:
             nu = jnp.array(default_params["nu"])
 
-        if isinstance(lengthscale, nnx.Variable):
+        if isinstance(lengthscale, paramax.AbstractUnwrappable):
             self.lengthscale = lengthscale
         else:
             self.lengthscale = PositiveReal(lengthscale)
 
         self.trainable_nu = trainable_nu
         if not trainable_nu:
-            self.nu = nu
-        elif isinstance(nu, nnx.Variable):
+            self.nu = paramax.non_trainable(jnp.asarray(paramax.unwrap(nu)))
+        elif isinstance(nu, paramax.AbstractUnwrappable):
             self.nu = nu
         else:
             self.nu = PositiveReal(nu)
 
-        if isinstance(variance, nnx.Variable):
+        if isinstance(variance, paramax.AbstractUnwrappable):
             self.variance = variance
         else:
             self.variance = NonNegativeReal(variance)

geometric_kernels/kernels/__init__.py

@@ -11,6 +11,9 @@
 from geometric_kernels.kernels.feature_map import MaternFeatureMapKernel
 from geometric_kernels.kernels.hodge_compositional import MaternHodgeCompositionalKernel
 from geometric_kernels.kernels.karhunen_loeve import MaternKarhunenLoeveKernel
+from geometric_kernels.kernels.karhunen_loeve_log_domain import (
+    MaternKarhunenLoeveLogDomain,
+)
 from geometric_kernels.kernels.matern_kernel import (
     MaternGeometricKernel,
     default_feature_map,

geometric_kernels/kernels/karhunen_loeve_log_domain.py

@@ -0,0 +1,81 @@
+"""Matérn Karhunen-Loève kernels with log-domain spectral normalization."""
+
+import lab as B
+import numpy as np
+from beartype.typing import Dict
+
+from geometric_kernels.kernels.karhunen_loeve import MaternKarhunenLoeveKernel
+from geometric_kernels.lab_extras import from_numpy, is_complex
+from geometric_kernels.utils.utils import _check_1_vector, _check_field_in_params
+
+
+class MaternKarhunenLoeveLogDomain(MaternKarhunenLoeveKernel):
+    """A discrete-spectrum Matérn kernel whose weights are computed in log space.
+
+    Eigenfunctions provide independent log multiplicities, allowing the
+    normalizer to be evaluated even when linear multiplicities overflow.
+    Eigenfunctions with ``supports_log_domain`` can also evaluate kernel
+    matrices without forming linear per-eigenfunction weights.
+    """
+
+    @staticmethod
+    def log_spectrum(s, nu, lengthscale, dimension):
+        """Evaluate the log Matérn spectrum without forming linear weights."""
+        _check_1_vector(lengthscale, "lengthscale")
+        _check_1_vector(nu, "nu")
+        s = B.cast(B.dtype(lengthscale), s)
+        safe_nu = B.where(nu == np.inf, B.ones(lengthscale), nu)
+        safe_lengthscale = B.where(nu == np.inf, B.ones(lengthscale), lengthscale)
+        finite = -(safe_nu + dimension / 2.0) * B.log(
+            2.0 * safe_nu / safe_lengthscale**2 + s
+        )
+        infinite = -(lengthscale**2) * s / 2.0
+        return B.where(nu == np.inf, infinite, finite)
+
+    @staticmethod
+    def spectrum(s, nu, lengthscale, dimension):
+        return B.exp(
+            MaternKarhunenLoeveLogDomain.log_spectrum(s, nu, lengthscale, dimension)
+        )
+
+    def _log_weights(self, params):
+        _check_field_in_params(params, "lengthscale")
+        _check_field_in_params(params, "nu")
+        log_spectrum = self.log_spectrum(
+            self.eigenvalues_laplacian,
+            params["nu"],
+            params["lengthscale"],
+            self.space.dimension,
+        )
+        log_multiplicities = B.cast(
+            B.dtype(log_spectrum),
+            from_numpy(
+                log_spectrum, self.eigenfunctions.log_num_eigenfunctions_per_level
+            ),
+        )[:, None]
+        log_levels = log_spectrum + log_multiplicities
+        return log_spectrum, log_levels, B.logsumexp(log_levels)
+
+    def log_eigenvalues(self, params: Dict[str, B.Numeric]) -> B.Numeric:
+        """Return per-eigenfunction log weights, shape [L, 1]."""
+        log_spectrum, _, log_normalizer = self._log_weights(params)
+        return log_spectrum - log_normalizer if self.normalize else log_spectrum
+
+    def eigenvalues(self, params: Dict[str, B.Numeric]) -> B.Numeric:
+        return B.exp(self.log_eigenvalues(params))
+
+    def K(self, params, X, X2=None, **kwargs):
+        if not self.eigenfunctions.supports_log_domain:
+            return super().K(params, X, X2, **kwargs)
+        result = self.eigenfunctions.weighted_outerproduct_log(
+            self.log_eigenvalues(params), X, X2, **kwargs
+        )
+        return B.real(result) if is_complex(result) else result
+
+    def K_diag(self, params, X, **kwargs):
+        if not self.eigenfunctions.supports_log_domain:
+            return super().K_diag(params, X, **kwargs)
+        result = self.eigenfunctions.weighted_outerproduct_diag_log(
+            self.log_eigenvalues(params), X, **kwargs
+        )
+        return B.real(result) if is_complex(result) else result