diff --git a/CHANGELOG.md b/CHANGELOG.md index 911fd3019..7b3482e39 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Fixed + +- **`PoweredExponential.power` is now constrained to $(0, 2]$ + ([#813](https://github.com/QuantClimate/GPJax/issues/813)).** The kernel is + positive definite in every dimension only in this interval, but any value was + accepted, and `power=3.0` gave a Gram matrix with negative eigenvalues. A + float in $(0, 2)$ is now wrapped as a trainable `SigmoidBounded` parameter, + so `fit` learns the power and cannot move it outside the interval. Before + this change a float power was a plain Python float, which `fit` did not + train. `power=2.0` (the RBF case) stays fixed, and a value outside the + interval raises `ValueError`. The type of the `power` leaf changes, so code + that reads it should use `val(kernel.power)`. + ### Changed - Allow JAX and JAXlib 0.11 in downstream environments by removing the diff --git a/gpjax/kernels/stationary/powered_exponential.py b/gpjax/kernels/stationary/powered_exponential.py index f2d72d1fe..d7d6e67dd 100644 --- a/gpjax/kernels/stationary/powered_exponential.py +++ b/gpjax/kernels/stationary/powered_exponential.py @@ -18,6 +18,7 @@ import beartype.typing as tp import jax.numpy as jnp from jaxtyping import Float +import paramax from paramax import AbstractUnwrappable from gpjax.kernels.base import val @@ -27,6 +28,7 @@ ) from gpjax.kernels.stationary.base import StationaryKernel from gpjax.kernels.stationary.utils import euclidean_distance +from gpjax.parameters import SigmoidBounded from gpjax.typing import ( Array, ScalarArray, @@ -50,6 +52,13 @@ class PoweredExponential(StationaryKernel): See Diggle and Ribeiro (2007) - "Model-based Geostatistics". and https://en.wikipedia.org/wiki/Generalized_normal_distribution#Symmetric_version + + The kernel is positive definite in every dimension only for + $0 < \kappa \le 2$. A float `power` in $(0, 2)$ becomes a trainable + parameter that the optimiser keeps inside $(0, 2)$. `power=2.0` gives the + RBF kernel and is held fixed, because a bounded parameter cannot sit on its + bound; to learn the power, start it inside the interval, for example at + 1.9. A parameter that you pass yourself must keep the power in $(0, 2]$. """ name: ClassVar[str] = "Powered Exponential" @@ -64,7 +73,7 @@ def __init__( n_dims: tp.Union[int, None] = None, compute_engine: AbstractKernelComputation = DenseKernelComputation(), ): - """Initializes the kernel. + r"""Initializes the kernel. Args: active_dims: the indices of the input dimensions that the kernel operates on. @@ -73,13 +82,16 @@ def __init__( used for all input dimensions. If an array with length > 1, the kernel is anisotropic, meaning that a different lengthscale is used for each input. variance: the variance of the kernel σ. - power: the power of the kernel κ. + power: the power of the kernel κ, in $(0, 2]$. n_dims: the number of input dimensions. If `lengthscale` is an array, this argument is ignored. compute_engine: the computation engine that the kernel uses to compute the covariance matrix. + + Raises: + ValueError: if the power is not in $(0, 2]$. """ - self.power = power + self.power = _wrap_power(power) super().__init__(active_dims, lengthscale, variance, n_dims, compute_engine) @@ -91,3 +103,18 @@ def __call__( power_val = val(self.power) K = val(self.variance) * jnp.exp(-(euclidean_distance(x, y) ** power_val)) return K.squeeze() + + +def _wrap_power(power: tp.Any) -> tp.Any: + value = float(val(power)) + if not 0.0 < value <= 2.0: + raise ValueError( + "Expected `power` in (0, 2], where the powered exponential kernel is " + f"positive definite. Got {value}." + ) + if isinstance(power, AbstractUnwrappable): + return power + power = jnp.asarray(power, dtype=float) + if value == 2.0: + return paramax.non_trainable(power) + return SigmoidBounded(power, low=0.0, high=2.0) diff --git a/tests/test_kernels/test_stationary.py b/tests/test_kernels/test_stationary.py index c36c41ea9..d394e8c44 100644 --- a/tests/test_kernels/test_stationary.py +++ b/tests/test_kernels/test_stationary.py @@ -32,11 +32,14 @@ from gpjax.parameters import ( NonNegativeReal, PositiveReal, + SigmoidBounded, + val, ) import jax from jax import config import jax.numpy as jnp import lineax as lx +import paramax from paramax import AbstractUnwrappable import pytest @@ -254,3 +257,46 @@ def test_name_is_not_a_constructor_argument(kernel: type[StationaryKernel]): assert "name" not in [f.name for f in dataclasses.fields(kernel)] with pytest.raises(TypeError): kernel(name="renamed") + + +# --------------------------------------------------------------------------- +# PoweredExponential power is constrained to (0, 2] (#813). +# --------------------------------------------------------------------------- + + +def test_powered_exponential_power_is_trainable_and_bounded(): + kernel = PoweredExponential(power=1.5) + assert isinstance(kernel.power, SigmoidBounded) + assert jnp.allclose(val(kernel.power), 1.5) + # A large step in the unconstrained space cannot leave the interval. + stepped = jax.tree_util.tree_map(lambda leaf: leaf + 100.0, kernel) + assert 0.0 < float(val(stepped.power)) <= 2.0 + + +def test_powered_exponential_power_two_is_fixed(): + kernel = PoweredExponential(power=2.0) + assert isinstance(kernel.power, paramax.NonTrainable) + assert jnp.allclose(val(kernel.power), 2.0) + + +@pytest.mark.parametrize("power", [-1.0, 0.0, 2.5, 3.0]) +def test_powered_exponential_rejects_power_outside_the_interval(power): + with pytest.raises(ValueError, match=r"\(0, 2\]"): + PoweredExponential(power=power) + + +def test_powered_exponential_rejects_an_invalid_parameter(): + with pytest.raises(ValueError, match=r"\(0, 2\]"): + PoweredExponential(power=PositiveReal(3.0)) + + +def test_powered_exponential_accepts_a_valid_parameter(): + power = PositiveReal(1.2) + assert PoweredExponential(power=power).power is power + + +@pytest.mark.parametrize("power", [0.3, 1.0, 1.7, 2.0]) +def test_powered_exponential_is_positive_definite(power): + X = jnp.linspace(0.0, 5.0, 200)[:, None] + gram = PoweredExponential(power=power).gram(X).as_matrix() + assert float(jnp.linalg.eigvalsh(gram).min()) > -1e-8