fix(kernels): parameterise the spectral measure by lengthscale and dimension

CHANGELOG.md

@@ -6,6 +6,34 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 
+## [Unreleased]
+
+### Fixed
+
+- **`StationaryKernel.spectral_density`**: now returns the correctly
+  parameterised spectral measure. It previously returned a *standardised*
+  distribution (`Normal(0, 1)` for RBF, `StudentT(2ν, 0, 1)` for Matérn) that
+  ignored the lengthscale and was hard-coded to one dimension, so
+  `kernel.spectral_density.log_prob(ω)` gave the same curve for every ℓ
+  ([#612](https://github.com/JaxGaussianProcesses/GPJax/issues/612)). The
+  measure is now `D`-dimensional and carries `diag(ℓ)⁻¹` as its scale (ARD
+  lengthscales included), satisfying Bochner's theorem
+  `k(τ) = σ²·E_p(ω)[exp(i ωᵀτ)]`.
+
+### Changed
+
+- **`StationaryKernel.spectral_density`** return type is now
+  `MultivariateNormal` / `MultivariateStudentT` with `event_shape == (D,)`,
+  where it was previously the univariate `Normal` / `StudentT`. Code calling
+  `.sample(key, (M, D))` should now call `.sample(key, (M,))`.
+- **`kernels.approximations.RFF`** with an explicitly supplied `frequencies=`
+  argument now treats those values as the spectral frequencies ω directly.
+  They were previously divided by the lengthscale inside
+  `BasisFunctionComputation.compute_features`, silently rescaling user-supplied
+  frequencies. RFF Gram/cross-covariance values with *sampled* frequencies are
+  bit-identical to `0.17.0` — the lengthscale simply moved from the feature map
+  into the measure it is drawn from.
+
 ## [0.17.0] — 2026-07-04
 
 ### Fixed

gpjax/kernels/approximations/rff.py

@@ -1,10 +1,8 @@
 """Compute Random Fourier Feature (RFF) kernel approximations."""
 
 import beartype.typing as tp
-import jax.numpy as jnp
 import jax.random as jr
 from jaxtyping import Float
-import numpyro.distributions as npd
 
 from gpjax.kernels.base import AbstractKernel
 from gpjax.kernels.computations import BasisFunctionComputation
@@ -65,20 +63,12 @@ def __init__(
                     "Expected the number of dimensions to be specified for the base kernel. "
                     "Please specify the n_dims argument for the base kernel."
                 )
-            sd = base_kernel.spectral_density
-            if isinstance(sd, npd.StudentT):
-                # Isotropic Matérn: one shared inverse-χ² per row couples the
-                # dimensions. iid-per-dim StudentT would give a tensor-product
-                # Matérn instead. MVT(df, scale_tril=I) has StudentT(df)
-                # marginals, so 1-D output is bit-identical to before.
-                mvt = npd.MultivariateStudentT(
-                    df=sd.df, loc=jnp.zeros(n_dims), scale_tril=jnp.eye(n_dims)
-                )
-                frequencies = mvt.sample(key, (num_basis_fns,))  # [M, D], coupled
-            else:
-                # Normal spectral density (RBF): a product of Gaussians is
-                # already isotropic, so iid-per-dim is correct.
-                frequencies = sd.sample(key=key, sample_shape=(num_basis_fns, n_dims))
+            # The spectral measure is a D-dimensional distribution already
+            # carrying the lengthscale as its inverse scale, so the draws are
+            # the frequencies ω themselves — no further rescaling downstream.
+            frequencies = base_kernel.spectral_density.sample(
+                key, (num_basis_fns,)
+            )  # [M, D]
 
         self.base_kernel = base_kernel
         self.num_basis_fns = num_basis_fns

gpjax/kernels/computations/basis_functions.py

@@ -56,9 +56,9 @@ def compute_features(
         Returns:
             A matrix of shape $N \times L$ representing the random fourier features where $L = 2M$.
         """
-        frequencies = kernel.frequencies
-        scaling_factor = _val(kernel.base_kernel.lengthscale)
-        z = jnp.matmul(x, (frequencies / scaling_factor).T)
+        # `frequencies` are the spectral frequencies ω, drawn from a measure
+        # that already carries the lengthscale as its inverse scale.
+        z = jnp.matmul(x, kernel.frequencies.T)
         z = jnp.concatenate([jnp.cos(z), jnp.sin(z)], axis=-1)
         return z
 

gpjax/kernels/stationary/base.py

@@ -20,7 +20,7 @@
 import numpyro.distributions as npd
 from paramax import AbstractUnwrappable
 
-from gpjax.kernels.base import AbstractKernel, _compute_base_init
+from gpjax.kernels.base import AbstractKernel, _compute_base_init, _val
 from gpjax.kernels.computations import (
     AbstractKernelComputation,
     DenseKernelComputation,
@@ -95,12 +95,48 @@ def __init__(
         self.n_dims = n_dims
         self.compute_engine = compute_engine
 
+    def _spectral_scale_tril(self) -> Float[Array, "D D"]:
+        r"""The scale matrix $\mathrm{diag}(1/\ell)$ of the spectral measure.
+
+        The spectral measure of a stationary kernel scales inversely with the
+        lengthscale: short lengthscales give wide spectra. Broadcasts a scalar
+        (isotropic) lengthscale across all $D$ dimensions and uses an ARD
+        lengthscale vector elementwise.
+        """
+        if self.n_dims is None:
+            raise ValueError(
+                f"Expected the number of dimensions to be specified for {self.name} "
+                "in order to construct its spectral measure. Please specify the "
+                "n_dims argument for the kernel."
+            )
+        return jnp.diag(jnp.ones(self.n_dims) / _val(self.lengthscale))
+
     @property
-    def spectral_density(self) -> npd.Normal | npd.StudentT:
-        r"""The spectral density of the kernel.
+    def spectral_density(self) -> npd.MultivariateNormal | npd.MultivariateStudentT:
+        r"""The normalised spectral measure $p(\boldsymbol{\omega})$ of the kernel.
+
+        By Bochner's theorem, a stationary kernel is the Fourier transform of a
+        finite measure. This property returns that measure *normalised to a
+        probability distribution over* $\mathbb{R}^D$, so that
+
+        $$
+        k(\boldsymbol{\tau}) = \sigma^2 \,
+            \mathbb{E}_{p(\boldsymbol{\omega})}
+            \big[e^{i \boldsymbol{\omega}^\top \boldsymbol{\tau}}\big].
+        $$
+
+        The measure depends on the lengthscale $\ell$ (as an inverse scale) but
+        **not** on the variance $\sigma^2$: the variance is the measure's total
+        mass, which normalisation divides out, and it re-enters as the explicit
+        prefactor above. The unnormalised spectral density of Rasmussen &
+        Williams (2006, §4.2.1) is recovered as
+        $S(\boldsymbol{\omega}) = \sigma^2 (2\pi)^D p(\boldsymbol{\omega})$
+        under the convention
+        $k(\boldsymbol{\tau}) = (2\pi)^{-D}\int S(\boldsymbol{\omega})
+        e^{i\boldsymbol{\omega}^\top\boldsymbol{\tau}}\,d\boldsymbol{\omega}$.
 
         Returns:
-            Callable[[Float[Array, "D"]], Float[Array, "D"]]: The spectral density function.
+            The spectral measure as a $D$-dimensional numpyro distribution.
         """
         raise NotImplementedError(
             f"Kernel {self.name} does not have a spectral density."

gpjax/kernels/stationary/matern12.py

@@ -49,5 +49,9 @@ def __call__(self, x: Float[Array, " D"], y: Float[Array, " D"]) -> ScalarFloat:
         return K.squeeze()
 
     @property
-    def spectral_density(self) -> npd.StudentT:
-        return build_student_t_distribution(nu=1)
+    def spectral_density(self) -> npd.MultivariateStudentT:
+        r"""The spectral measure of the Matérn-1/2 kernel: a multivariate
+        Student's t with 1 degree of freedom and scale $\mathrm{diag}(\ell)^{-1}$."""
+        return build_student_t_distribution(
+            nu=1, scale_tril=self._spectral_scale_tril()
+        )

gpjax/kernels/stationary/matern32.py

@@ -55,5 +55,9 @@ def __call__(
         return K.squeeze()
 
     @property
-    def spectral_density(self) -> npd.StudentT:
-        return build_student_t_distribution(nu=3)
+    def spectral_density(self) -> npd.MultivariateStudentT:
+        r"""The spectral measure of the Matérn-3/2 kernel: a multivariate
+        Student's t with 3 degrees of freedom and scale $\mathrm{diag}(\ell)^{-1}$."""
+        return build_student_t_distribution(
+            nu=3, scale_tril=self._spectral_scale_tril()
+        )

gpjax/kernels/stationary/matern52.py

@@ -54,5 +54,9 @@ def __call__(
         return K.squeeze()
 
     @property
-    def spectral_density(self) -> npd.StudentT:
-        return build_student_t_distribution(nu=5)
+    def spectral_density(self) -> npd.MultivariateStudentT:
+        r"""The spectral measure of the Matérn-5/2 kernel: a multivariate
+        Student's t with 5 degrees of freedom and scale $\mathrm{diag}(\ell)^{-1}$."""
+        return build_student_t_distribution(
+            nu=5, scale_tril=self._spectral_scale_tril()
+        )

gpjax/kernels/stationary/rbf.py

@@ -45,5 +45,9 @@ def __call__(self, x: Float[Array, " D"], y: Float[Array, " D"]) -> ScalarFloat:
         return K.squeeze()
 
     @property
-    def spectral_density(self) -> npd.Normal:
-        return npd.Normal(0.0, 1.0)
+    def spectral_density(self) -> npd.MultivariateNormal:
+        r"""The spectral measure $\mathcal{N}(\boldsymbol{0}, \mathrm{diag}(\ell)^{-2})$."""
+        scale_tril = self._spectral_scale_tril()
+        return npd.MultivariateNormal(
+            jnp.zeros(scale_tril.shape[0]), scale_tril=scale_tril
+        )

gpjax/kernels/stationary/utils.py

@@ -22,21 +22,29 @@
 )
 
 
-def build_student_t_distribution(nu: int) -> npd.StudentT:
-    r"""Build a Student's t distribution with a fixed smoothness parameter.
+def build_student_t_distribution(
+    nu: int, scale_tril: Float[Array, "D D"]
+) -> npd.MultivariateStudentT:
+    r"""Build the spectral measure of a Matérn kernel.
 
-    For a fixed half-integer smoothness parameter, compute the spectral density of a
-    Matérn kernel; a Student's t distribution.
+    The Matérn kernel with smoothness $\nu$ has spectral density proportional to
+    $(2\nu/\ell^2 + \lVert\omega\rVert^2)^{-(\nu + D/2)}$, which normalises to a
+    multivariate Student's t measure with $2\nu$ degrees of freedom and scale
+    matrix $\mathrm{diag}(\ell)^{-1}$.
 
     Args:
-        nu (int): The smoothness parameter of the Matérn kernel.
+        nu (int): Twice the smoothness parameter of the Matérn kernel, i.e. 1, 3
+            or 5 for the Matérn-1/2, -3/2 and -5/2 kernels respectively.
+        scale_tril (Float[Array, "D D"]): The scale matrix $\mathrm{diag}(\ell)^{-1}$
+            of the measure, as returned by `StationaryKernel._spectral_scale_tril`.
 
     Returns
     -------
-        tfp.Distribution: A Student's t distribution with the same smoothness parameter.
+        npd.MultivariateStudentT: The spectral measure over $\mathbb{R}^D$.
     """
-    dist = npd.StudentT(df=nu, loc=0.0, scale=1.0)
-    return dist
+    return npd.MultivariateStudentT(
+        df=nu, loc=jnp.zeros(scale_tril.shape[0]), scale_tril=scale_tril
+    )
 
 
 def squared_distance(x: Float[Array, " D"], y: Float[Array, " D"]) -> ScalarFloat:

tests/test_kernels/test_approximations.py

@@ -18,6 +18,7 @@
 import jax.numpy as jnp
 import jax.random as jr
 import lineax as lx
+import numpyro.distributions as npd
 import pytest
 
 config.update("jax_enable_x64", True)
@@ -185,11 +186,22 @@ def test_matern_rff_gram_matches_isotropic(n_dims):
     assert rel_frobenius < 0.02
 
 
-def test_matern_rff_frequencies_1d_unchanged():
-    """1-D behaviour is preserved bit-identically: numpyro's 1-D MVT consumes
-    the PRNG stream exactly as the univariate StudentT it replaces."""
-    base_kernel = Matern32(active_dims=[0])
+@pytest.mark.parametrize("n_dims", [1, 3])
+@pytest.mark.parametrize("lengthscale", [0.5, 1.0, 2.0])
+def test_rff_effective_frequencies_unchanged(n_dims: int, lengthscale: float):
+    """Folding the lengthscale into the spectral measure is bit-identical.
+
+    Previously the measure was standardised and `compute_features` divided the
+    draws by ℓ. Now the measure carries diag(ℓ)⁻¹ directly and no rescaling
+    happens downstream. Both routes must give the same effective frequencies.
+    """
+    base_kernel = Matern32(n_dims=n_dims, lengthscale=lengthscale)
     approx = RFF(base_kernel=base_kernel, num_basis_fns=64, key=jr.key(5))
-    # Reference: the pre-fix univariate draw for the same key/shape.
-    reference = base_kernel.spectral_density.sample(key=jr.key(5), sample_shape=(64, 1))
-    assert jnp.allclose(approx.frequencies, reference)
+
+    # Reference: the pre-fix route — standardised draw, then divide by ℓ.
+    standardised = npd.MultivariateStudentT(
+        df=3, loc=jnp.zeros(n_dims), scale_tril=jnp.eye(n_dims)
+    )
+    reference = standardised.sample(jr.key(5), (64,)) / lengthscale
+
+    assert jnp.array_equal(approx.frequencies, reference)