diff --git a/CHANGELOG.md b/CHANGELOG.md index eb1ecde57..8dbcc3646 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/gpjax/kernels/approximations/rff.py b/gpjax/kernels/approximations/rff.py index bed616019..d9c9595f4 100644 --- a/gpjax/kernels/approximations/rff.py +++ b/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 diff --git a/gpjax/kernels/computations/basis_functions.py b/gpjax/kernels/computations/basis_functions.py index e205848fc..6f6d1dcc9 100644 --- a/gpjax/kernels/computations/basis_functions.py +++ b/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 diff --git a/gpjax/kernels/stationary/base.py b/gpjax/kernels/stationary/base.py index a5c37af79..f5b8c4423 100644 --- a/gpjax/kernels/stationary/base.py +++ b/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." diff --git a/gpjax/kernels/stationary/matern12.py b/gpjax/kernels/stationary/matern12.py index 2081eb800..29ea7a484 100644 --- a/gpjax/kernels/stationary/matern12.py +++ b/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() + ) diff --git a/gpjax/kernels/stationary/matern32.py b/gpjax/kernels/stationary/matern32.py index 6914ff766..8f1e575d4 100644 --- a/gpjax/kernels/stationary/matern32.py +++ b/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() + ) diff --git a/gpjax/kernels/stationary/matern52.py b/gpjax/kernels/stationary/matern52.py index 6c82c5e5b..b98ac59f4 100644 --- a/gpjax/kernels/stationary/matern52.py +++ b/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() + ) diff --git a/gpjax/kernels/stationary/rbf.py b/gpjax/kernels/stationary/rbf.py index bc87074c9..791dd57e0 100644 --- a/gpjax/kernels/stationary/rbf.py +++ b/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 + ) diff --git a/gpjax/kernels/stationary/utils.py b/gpjax/kernels/stationary/utils.py index bbe7b0a7d..8651f023e 100644 --- a/gpjax/kernels/stationary/utils.py +++ b/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: diff --git a/tests/test_kernels/test_approximations.py b/tests/test_kernels/test_approximations.py index 8bd649e98..89d38c474 100644 --- a/tests/test_kernels/test_approximations.py +++ b/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) diff --git a/tests/test_kernels/test_spectral_density.py b/tests/test_kernels/test_spectral_density.py new file mode 100644 index 000000000..d99647a1d --- /dev/null +++ b/tests/test_kernels/test_spectral_density.py @@ -0,0 +1,140 @@ +# Copyright 2022 The thomaspinder Contributors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for the spectral measures of stationary kernels (issue #612).""" + +from itertools import pairwise + +from gpjax.kernels.stationary import ( + RBF, + Matern12, + Matern32, + Matern52, +) +from gpjax.kernels.stationary.base import StationaryKernel +import jax +from jax import config +import jax.numpy as jnp +import jax.random as jr +import pytest + +# Enable Float64 for more stable matrix inversions. +config.update("jax_enable_x64", True) + +SPECTRAL_KERNELS = [RBF, Matern12, Matern32, Matern52] + + +@pytest.mark.parametrize("kernel", SPECTRAL_KERNELS) +@pytest.mark.parametrize("n_dims", [1, 2, 3]) +@pytest.mark.parametrize("lengthscale", [0.5, 1.0, 2.0]) +@pytest.mark.parametrize("variance", [1.0, 3.0]) +def test_spectral_measure_satisfies_bochner( + kernel: type[StationaryKernel], + n_dims: int, + lengthscale: float, + variance: float, +): + """Bochner's theorem: k(τ) = σ² · E_{p(ω)}[cos(ωᵀτ)]. + + This pins the spectral measure's parameterisation without reference to any + Fourier-transform normalisation convention. A measure that ignores the + lengthscale (issue #612) fails this for every ℓ != 1. + """ + base_kernel = kernel(n_dims=n_dims, lengthscale=lengthscale, variance=variance) + measure = base_kernel.spectral_density + + omega = measure.sample(jr.key(42), (200_000,)) + assert omega.shape == (200_000, n_dims) + + # A handful of separations at which to check the identity. + tau = jnp.array([[0.0], [0.3], [0.7], [1.5]]) * jnp.ones((1, n_dims)) + + monte_carlo = variance * jnp.mean(jnp.cos(omega @ tau.T), axis=0) + exact = jax.vmap(lambda t: base_kernel(jnp.zeros(n_dims), t))(tau) + + assert jnp.allclose(monte_carlo, exact, atol=2e-2) + + +@pytest.mark.parametrize("kernel", SPECTRAL_KERNELS) +def test_spectral_density_depends_on_lengthscale(kernel: type[StationaryKernel]): + """The reporter's reproducer: densities for different ℓ must not coincide. + + Previously every lengthscale returned the same standardised measure, so the + curves were superimposed. + """ + omega = jnp.linspace(-5.0, 5.0, 101).reshape(-1, 1) + + log_probs = [ + kernel(n_dims=1, lengthscale=ell).spectral_density.log_prob(omega) + for ell in (0.5, 1.0, 2.0) + ] + + for a, b in pairwise(log_probs): + assert not jnp.allclose(a, b) + + +@pytest.mark.parametrize("kernel", SPECTRAL_KERNELS) +def test_spectral_measure_is_multivariate(kernel: type[StationaryKernel]): + """The measure must live on R^D, not be hard-coded to 1-D.""" + n_dims = 3 + measure = kernel(n_dims=n_dims).spectral_density + + assert measure.event_shape == (n_dims,) + + omega = jnp.ones((7, n_dims)) + assert measure.log_prob(omega).shape == (7,) + + +@pytest.mark.parametrize("kernel", SPECTRAL_KERNELS) +def test_spectral_measure_scales_inversely_with_lengthscale( + kernel: type[StationaryKernel], +): + """A longer lengthscale must concentrate the measure near the origin.""" + short = kernel(n_dims=1, lengthscale=0.5).spectral_density + long = kernel(n_dims=1, lengthscale=4.0).spectral_density + + key = jr.key(0) + spread_short = jnp.mean(jnp.abs(short.sample(key, (20_000,)))) + spread_long = jnp.mean(jnp.abs(long.sample(key, (20_000,)))) + + assert spread_short > spread_long + + +@pytest.mark.parametrize("kernel", SPECTRAL_KERNELS) +def test_spectral_measure_is_anisotropic_under_ard(kernel: type[StationaryKernel]): + """An ARD lengthscale must give a per-dimension spectral scale.""" + base_kernel = kernel(lengthscale=jnp.array([0.5, 4.0])) + measure = base_kernel.spectral_density + + omega = measure.sample(jr.key(1), (20_000,)) + spread = jnp.mean(jnp.abs(omega), axis=0) + + # Dimension 0 has the shorter lengthscale, so the wider spectrum. + assert spread[0] > spread[1] + + +@pytest.mark.parametrize("kernel", SPECTRAL_KERNELS) +def test_spectral_measure_is_independent_of_variance( + kernel: type[StationaryKernel], +): + """σ² is the measure's total mass, not part of its shape. + + `spectral_density` returns a normalised probability measure, so the variance + must not appear in it; Bochner's theorem carries σ² as a separate factor. + """ + unit = kernel(n_dims=1, variance=1.0).spectral_density + scaled = kernel(n_dims=1, variance=5.0).spectral_density + + omega = jnp.linspace(-3.0, 3.0, 51).reshape(-1, 1) + assert jnp.allclose(unit.log_prob(omega), scaled.log_prob(omega)) diff --git a/tests/test_kernels/test_utils.py b/tests/test_kernels/test_utils.py index 05e1265ca..c8956938e 100644 --- a/tests/test_kernels/test_utils.py +++ b/tests/test_kernels/test_utils.py @@ -98,18 +98,22 @@ def test_euclidean_distance_same_point() -> None: @pytest.mark.parametrize("nu", [1, 3, 5, 10]) -def test_build_student_t_distribution(nu: int) -> None: - dist = build_student_t_distribution(nu) - assert isinstance(dist, npd.StudentT) +@pytest.mark.parametrize("n_dims", [1, 3]) +def test_build_student_t_distribution(nu: int, n_dims: int) -> None: + scale_tril = jnp.eye(n_dims) + dist = build_student_t_distribution(nu, scale_tril=scale_tril) + assert isinstance(dist, npd.MultivariateStudentT) assert dist.df == nu - assert dist.loc == 0.0 - assert dist.scale == 1.0 + assert dist.event_shape == (n_dims,) + assert jnp.allclose(dist.loc, jnp.zeros(n_dims)) + assert jnp.allclose(dist.scale_tril, scale_tril) -def test_student_t_is_sampleable() -> None: +@pytest.mark.parametrize("n_dims", [1, 3]) +def test_student_t_is_sampleable(n_dims: int) -> None: import jax.random as jr - dist = build_student_t_distribution(5) + dist = build_student_t_distribution(5, scale_tril=jnp.eye(n_dims)) samples = dist.sample(jr.key(0), (100,)) - assert samples.shape == (100,) + assert samples.shape == (100, n_dims) assert jnp.all(jnp.isfinite(samples))