Add changepoint kernel

gpjax/kernels/stationary/__init__.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
+from gpjax.kernels.stationary.changepoint import ChangePoint
 from gpjax.kernels.stationary.matern12 import Matern12
 from gpjax.kernels.stationary.matern32 import Matern32
 from gpjax.kernels.stationary.matern52 import Matern52
@@ -23,6 +24,7 @@
 from gpjax.kernels.stationary.white import White
 
 __all__ = [
+    "ChangePoint",
     "Matern12",
     "Matern32",
     "Matern52",

gpjax/kernels/stationary/changepoint.py

@@ -0,0 +1,115 @@
+# Copyright 2023 The JaxGaussianProcesses 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.
+# ==============================================================================
+
+
+from dataclasses import dataclass
+
+from beartype.typing import (
+    Callable,
+    List,
+)
+import jax
+import jax.numpy as jnp
+from jaxtyping import Float
+import tensorflow_probability.substrates.jax.bijectors as tfb
+
+from gpjax.base import (
+    param_field,
+    static_field,
+)
+from gpjax.kernels.base import AbstractKernel
+from gpjax.typing import (
+    Array,
+    ScalarFloat,
+)
+
+
+@dataclass
+class ChangePoint(AbstractKernel):
+    r"""A change point kernel
+    self.kernels: A list of exactly two kernels that will be switched.
+    self.tswitch: The point at which to change to a different kernel.
+        for example: if x and y are both less than tswitch, then you would use kernels[0]
+                     if x and y are both greater than or equal to tswitch, then you would use kernels[1]
+                     otherwise return cross-covariance of 0
+    """
+
+    kernels: List[AbstractKernel] = None
+    operator: Callable = static_field(None)
+    tswitch: ScalarFloat = param_field(
+        jnp.array(1.0), bijector=tfb.Identity(), trainable=False
+    )
+
+    def __post_init__(self):
+        # Add kernels to a list, flattening out instances of this class therein, as in GPFlow kernels.
+        kernels_list: List[AbstractKernel] = []
+
+        for kernel in self.kernels:
+            if not isinstance(kernel, AbstractKernel):
+                raise TypeError("can only combine Kernel instances")  # pragma: no cover
+
+            if isinstance(kernel, self.__class__):
+                kernels_list.extend(kernel.kernels)
+            else:
+                kernels_list.append(kernel)
+
+        self.kernels = kernels_list
+
+    def __call__(
+        self,
+        x: Float[Array, " D"],
+        y: Float[Array, " D"],
+    ) -> ScalarFloat:
+        r"""Evaluate the kernel on a pair of inputs.
+
+        Args:
+            x (Float[Array, " D"]): The left hand input of the kernel function.
+            y (Float[Array, " D"]): The right hand input of the kernel function.
+
+        Returns
+        -------
+            ScalarFloat: The evaluated kernel function at the supplied inputs.
+        """
+
+        def get_function_index(x, y, tswitch=1):
+            r"""
+            Specify four possible indices given x, y, and tswitch.
+
+            Args:
+                    x: Left hand argument of kernel function's call
+                    y: Right hand argument of kernel function's call
+                    tswitch: point at which to change to a different kernel
+            """
+            # Four possible indexes: 0, 1, 2, 3
+            indx = 3  # if indx doesn't get set to 0, 1, or 2, then default 3
+            # lessthan means that at tswitch, you are already switched
+            cond1 = jnp.less(x, tswitch)
+            cond2 = jnp.less(y, tswitch)
+            indx = jnp.where(jnp.logical_and(cond1, cond2), 0, indx)
+            indx = jnp.where(jnp.logical_and(jnp.invert(cond1), cond2), 1, indx)
+            indx = jnp.where(jnp.logical_and(cond1, jnp.invert(cond2)), 2, indx)
+            return indx.squeeze().astype("uint8")
+
+        def k_zero(x, y):
+            r"""Return 0 covariance"""
+            out = jnp.float64(0)
+            return out.squeeze()
+
+        indx = get_function_index(x, y, tswitch=self.tswitch)
+
+        flst = [self.kernels[0].__call__, k_zero, k_zero, self.kernels[1].__call__]
+        K = jax.lax.switch(indx, flst, x, y)
+
+        return K.squeeze()

tests/test_kernels/test_stationary.py

@@ -34,6 +34,7 @@
 )
 from gpjax.kernels.stationary import (
     RBF,
+    ChangePoint,
     Matern12,
     Matern32,
     Matern52,
@@ -241,6 +242,44 @@ class TestRationalQuadratic(BaseTestKernel):
     default_compute_engine = DenseKernelComputation()
 
 
+@pytest.mark.parametrize(
+    "kernel",
+    [RBF, RationalQuadratic, Matern12, Matern32, Matern52],
+)
+@pytest.mark.parametrize("tswitch_param", [0.1, 0.25, 0.5, 0.75, 0.9])
+def test_changepoint_kernel(kernel: AbstractKernel, tswitch_param: float) -> None:
+    n_kerns = 2
+    # Create inputs
+    n = 20
+    x = jnp.linspace(0.0, 1.0, num=n).reshape(-1, 1)
+
+    # Create list of kernels
+    kernels = [kernel() for _ in range(n_kerns)]
+
+    # Create combination kernel
+    combination_kernel = ChangePoint(kernels=kernels, tswitch=tswitch_param)
+
+    # Check params are a list of dictionaries
+    assert combination_kernel.kernels == kernels
+
+    # Check combination kernel set
+    assert len(combination_kernel.kernels) == n_kerns
+    assert isinstance(combination_kernel.kernels, list)
+    assert isinstance(combination_kernel.kernels[0], AbstractKernel)
+
+    # Compute gram matrix
+    Kxx = combination_kernel.gram(x)
+
+    # Check shapes
+    assert Kxx.shape[0] == Kxx.shape[1]
+    assert Kxx.shape[1] == n
+
+    # Check positive definiteness
+    jitter = 1e-6
+    eigen_values = jnp.linalg.eigvalsh(Kxx.to_dense() + jnp.eye(n) * jitter)
+    assert (eigen_values > 0).all()
+
+
 @pytest.mark.parametrize("smoothness", [1, 2, 3])
 def test_build_studentt_dist(smoothness: int) -> None:
     dist = build_student_t_distribution(smoothness)