Implementing a diagonal kernel

[zwei-beiner]

Hi,

I am trying to implement a diagonal noise kernel (equivalent to scikit-learn's WhiteKernel). (I know that there is the diag option in tinygp but I am trying to create a block-diagonal covariance matrix for a multi-output GP, for which each block can have different diagonal noise, so I have to implement this as a kernel.)

However, I get strange results at the edges of the predictions:

• The error bars shrink to a small value.
• The mean prediction goes almost exactly through the training point at the edge:

I have marked the strange predictions with the red arrows:

I am trying to debug this behaviour, but I am not sure how. Any help would be appreaciated!

---

Code:

import dataclasses

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

import os
from pathlib import Path

import functools
from tinygp import kernels, GaussianProcess
import jax
import jax.numpy as jnp
import jaxopt

jax.config.update("jax_enable_x64", True)

@dataclasses.dataclass
class WhiteKernel(kernels.Kernel):
    noise: jnp.ndarray

    def true_fun(self):
        return self.noise

    def false_fun(self):
        return 0.

    def evaluate(self, X1, X2):
        # Should depend on matrix index, but cannot do this, so need to calculate from values instead
        return jax.lax.cond(X1 == X2, self.true_fun, self.false_fun)

# Toy data
rng = np.random.default_rng(seed=0)
t_train = np.sort(rng.uniform(-5, 5, size=20))
y_train = np.sin(t_train) + rng.normal(0, 0.2, size=t_train.shape)

fig: plt.Figure
ax: plt.Axes
fig, ax = plt.subplots(figsize=(9, 6))
ax.plot(t_train, y_train, '.', label='Data: sin(x)+noise')
ax.set_xlabel('t')
ax.set_ylabel('y')

def build_gp(params):
    kernel = kernels.Matern32(jnp.exp(params["log_gp_scale"])) + WhiteKernel(jnp.exp(params["log_gp_diag"]))
    return GaussianProcess(kernel, t_train)

@jax.jit
def loss(params):
    gp = build_gp(params)
    return -gp.log_probability(y_train)

theta_init = {
    "log_gp_scale": np.float64(1.6),
    "log_gp_diag": np.log(0.2),
}

solver = jaxopt.ScipyMinimize(fun=loss)
ML_II_params, info = solver.run(theta_init)

t_test = np.linspace(np.min(t_train), np.max(t_train), num=100)
gp = build_gp(ML_II_params)
_, gp_conditioned = gp.condition(y_train, t_test)
y_test, y_test_std = gp_conditioned.loc, np.sqrt(gp_conditioned.variance)

ax.plot(t_test, y_test)
ax.plot(t_test, y_test, color='orange', label='ML-II')
ax.fill_between(t_test, y_test - y_test_std, y_test + y_test_std, alpha=0.25, color='orange')

ax.legend()
plt.show()