fix: constrain PoweredExponential.power to (0, 2]

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

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)

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