jaxgaussianprocesses__gpjax METASEP

tests/test_utilities.py METASEP
import jax.numpy as jnp
import jax.random as jr
import pytest

from gpjax.utils import (
    I,
    as_constant,
    concat_dictionaries,
    dict_array_coercion,
    merge_dictionaries,
    sort_dictionary,
)


@pytest.mark.parametrize("n", [1, 10, 100])
def test_identity(n):
    identity = I(n)
    assert identity.shape == (n, n)
    assert (jnp.diag(identity) == jnp.ones(shape=(n,))).all()


def test_concat_dict():
    d1 = {"a": 1, "b": 2}
    d2 = {"c": 3, "d": 4}
    d = concat_dictionaries(d1, d2)
    assert list(d.keys()) == ["a", "b", "c", "d"]
    assert list(d.values()) == [1, 2, 3, 4]


def test_merge_dicts():
    d1 = {"a": 1, "b": 2}
    d2 = {"b": 3}
    d = merge_dictionaries(d1, d2)
    assert list(d.keys()) == ["a", "b"]
    assert list(d.values()) == [1, 3]


def test_sort_dict():
    unsorted = {"b": 1, "a": 2}
    sorted_dict = sort_dictionary(unsorted)
    assert list(sorted_dict.keys()) == ["a", "b"]
    assert list(sorted_dict.values()) == [2, 1]


def test_as_constant():
    base = {"a": 1, "b": 2, "c": 3}
    b1, s1 = as_constant(base, ["a"])
    b2, s2 = as_constant(base, ["a", "b"])
    assert list(b1.keys()) == ["b", "c"]
    assert list(s1.keys()) == ["a"]
    assert list(b2.keys()) == ["c"]
    assert list(s2.keys()) == ["a", "b"]


@pytest.mark.parametrize("d", [1, 2, 10])
def test_array_coercion(d):
    params = {
        "kernel": {
            "lengthscale": jnp.array([1.0] * d),
            "variance": jnp.array([1.0]),
        },
        "likelihood": {"obs_noise": jnp.array([1.0])},
        "mean_function": {},
    }
    dict_to_array, array_to_dict = dict_array_coercion(params)
    assert array_to_dict(dict_to_array(params)) == params
    assert isinstance(dict_to_array(params), list)
    assert isinstance(array_to_dict(dict_to_array(params)), dict)

tests/test_types.py METASEP
import jax.numpy as jnp
import pytest

from gpjax.types import Dataset, NoneType, verify_dataset


def test_nonetype():
    assert isinstance(None, NoneType)


@pytest.mark.parametrize("n", [1, 10, 100])
@pytest.mark.parametrize("outd", [1, 2, 10])
@pytest.mark.parametrize("ind", [1, 2, 10])
def test_dataset(n, outd, ind):
    x = jnp.ones((n, ind))
    y = jnp.ones((n, outd))
    d = Dataset(X=x, y=y)
    verify_dataset(d)
    assert d.n == n
    assert d.in_dim == ind
    assert d.out_dim == outd


@pytest.mark.parametrize("nx, ny", [(1, 2), (2, 1), (10, 5), (5, 10)])
def test_dataset_assertions(nx, ny):
    x = jnp.ones((nx, 1))
    y = jnp.ones((ny, 1))
    with pytest.raises(AssertionError):
        ds = Dataset(X=x, y=y)
        verify_dataset(ds)


def test_y_none():
    x = jnp.ones((10, 1))
    d = Dataset(X=x)
    verify_dataset(d)
    assert d.y is None

tests/test_parameters.py METASEP
import typing as tp
from argparse import _AppendConstAction
from copy import deepcopy

import jax.numpy as jnp
import pytest
from numpy.testing import assert_array_equal
from tensorflow_probability.substrates.jax import distributions as tfd

from gpjax.config import get_defaults
from gpjax.gps import Prior
from gpjax.kernels import RBF
from gpjax.likelihoods import Bernoulli, Gaussian
from gpjax.parameters import (  # build_all_transforms,
    build_transforms,
    copy_dict_structure,
    evaluate_priors,
    initialise,
    log_density,
    prior_checks,
    recursive_complete,
    recursive_items,
    structure_priors,
    transform,
)


#########################
# Test base functionality
#########################
@pytest.mark.parametrize("lik", [Gaussian])
def test_initialise(lik):
    posterior = Prior(kernel=RBF()) * lik(num_datapoints=10)
    params, _, _ = initialise(posterior)
    assert list(sorted(params.keys())) == [
        "kernel",
        "likelihood",
        "mean_function",
    ]


def test_non_conjugate_initialise():
    posterior = Prior(kernel=RBF()) * Bernoulli(num_datapoints=10)
    params, _, _ = initialise(posterior)
    assert list(sorted(params.keys())) == [
        "kernel",
        "latent",
        "likelihood",
        "mean_function",
    ]


#########################
# Test priors
#########################
@pytest.mark.parametrize("x", [-1.0, 0.0, 1.0])
def test_lpd(x):
    val = jnp.array(x)
    dist = tfd.Normal(loc=0.0, scale=1.0)
    lpd = log_density(val, dist)
    assert lpd is not None
    assert log_density(val, None) == 0.0


@pytest.mark.parametrize("lik", [Gaussian, Bernoulli])
def test_prior_template(lik):
    posterior = Prior(kernel=RBF()) * lik(num_datapoints=10)
    params, _, _ = initialise(posterior)
    prior_container = copy_dict_structure(params)
    for (
        k,
        v1,
        v2,
    ) in recursive_items(params, prior_container):
        assert v2 == None


@pytest.mark.parametrize("lik", [Gaussian, Bernoulli])
def test_recursive_complete(lik):
    posterior = Prior(kernel=RBF()) * lik(num_datapoints=10)
    params, _, _ = initialise(posterior)
    priors = {"kernel": {}}
    priors["kernel"]["lengthscale"] = tfd.HalfNormal(scale=2.0)
    container = copy_dict_structure(params)
    complete_priors = recursive_complete(container, priors)
    for (
        k,
        v1,
        v2,
    ) in recursive_items(params, complete_priors):
        if k == "lengthscale":
            assert isinstance(v2, tfd.HalfNormal)
        else:
            assert v2 == None


def test_prior_evaluation():
    """
    Test the regular setup that every parameter has a corresponding prior distribution attached to its unconstrained
    value.
    """
    params = {
        "kernel": {
            "lengthscale": jnp.array([1.0]),
            "variance": jnp.array([1.0]),
        },
        "likelihood": {"obs_noise": jnp.array([1.0])},
    }
    priors = {
        "kernel": {
            "lengthscale": tfd.Gamma(1.0, 1.0),
            "variance": tfd.Gamma(2.0, 2.0),
        },
        "likelihood": {"obs_noise": tfd.Gamma(3.0, 3.0)},
    }
    lpd = evaluate_priors(params, priors)
    assert pytest.approx(lpd) == -2.0110168


def test_none_prior():
    """
    Test that multiple dispatch is working in the case of no priors.
    """
    params = {
        "kernel": {
            "lengthscale": jnp.array([1.0]),
            "variance": jnp.array([1.0]),
        },
        "likelihood": {"obs_noise": jnp.array([1.0])},
    }
    priors = copy_dict_structure(params)
    lpd = evaluate_priors(params, priors)
    assert lpd == 0.0


def test_incomplete_priors():
    """
    Test the case where a user specifies priors for some, but not all, parameters.
    """
    params = {
        "kernel": {
            "lengthscale": jnp.array([1.0]),
            "variance": jnp.array([1.0]),
        },
        "likelihood": {"obs_noise": jnp.array([1.0])},
    }
    priors = {
        "kernel": {
            "lengthscale": tfd.Gamma(1.0, 1.0),
            "variance": tfd.Gamma(2.0, 2.0),
        },
    }
    container = copy_dict_structure(params)
    complete_priors = recursive_complete(container, priors)
    lpd = evaluate_priors(params, complete_priors)
    assert pytest.approx(lpd) == -1.6137061


@pytest.mark.parametrize("num_datapoints", [1, 10])
def test_checks(num_datapoints):
    incomplete_priors = {"lengthscale": jnp.array([1.0])}
    posterior = Prior(kernel=RBF()) * Bernoulli(num_datapoints=num_datapoints)
    priors = prior_checks(incomplete_priors)
    assert "latent" in priors.keys()
    assert "variance" not in priors.keys()
    assert isinstance(priors["latent"], tfd.Normal)


def test_structure_priors():
    posterior = Prior(kernel=RBF()) * Gaussian(num_datapoints=10)
    params, _, _ = initialise(posterior)
    priors = {
        "kernel": {
            "lengthscale": tfd.Gamma(1.0, 1.0),
            "variance": tfd.Gamma(2.0, 2.0),
        },
    }
    structured_priors = structure_priors(params, priors)

    def recursive_fn(d1, d2, fn: tp.Callable[[tp.Any], tp.Any]):
        for key, value in d1.items():
            if type(value) is dict:
                yield from recursive_fn(value, d2[key], fn)
            else:
                yield fn(key, key)

    for v in recursive_fn(params, structured_priors, lambda k1, k2: k1 == k2):
        assert v


def test_prior_checks():
    priors = {
        "kernel": {"lengthscale": None, "variance": None},
        "mean_function": {},
        "liklelihood": {"variance": None},
        "latent": None,
    }
    new_priors = prior_checks(priors)
    assert "latent" in new_priors.keys()
    assert new_priors["latent"].name == "Normal"

    priors = {
        "kernel": {"lengthscale": None, "variance": None},
        "mean_function": {},
        "liklelihood": {"variance": None},
    }
    new_priors = prior_checks(priors)
    assert "latent" in new_priors.keys()
    assert new_priors["latent"].name == "Normal"

    priors = {
        "kernel": {"lengthscale": None, "variance": None},
        "mean_function": {},
        "liklelihood": {"variance": None},
        "latent": tfd.StudentT(df=5, loc=0.0, scale=1.0),
    }
    with pytest.warns(UserWarning):
        new_priors = prior_checks(priors)
    assert "latent" in new_priors.keys()
    assert new_priors["latent"].name == "StudentT"


#########################
# Test transforms
#########################
@pytest.mark.parametrize("num_datapoints", [1, 10])
@pytest.mark.parametrize("likelihood", [Gaussian, Bernoulli])
def test_output(num_datapoints, likelihood):
    posterior = Prior(kernel=RBF()) * likelihood(num_datapoints=num_datapoints)
    params, constrainer, unconstrainer = initialise(posterior)

    assert isinstance(constrainer, dict)
    assert isinstance(unconstrainer, dict)
    for k, v1, v2 in recursive_items(constrainer, unconstrainer):
        assert isinstance(v1, tp.Callable)
        assert isinstance(v2, tp.Callable)

    unconstrained_params = transform(params, unconstrainer)
    assert (
        unconstrained_params["kernel"]["lengthscale"] != params["kernel"]["lengthscale"]
    )
    backconstrained_params = transform(unconstrained_params, constrainer)
    for k, v1, v2 in recursive_items(params, unconstrained_params):
        assert v1.dtype == v2.dtype

    for k, v1, v2 in recursive_items(params, backconstrained_params):
        assert all(v1 == v2)

    augmented_params = params
    augmented_params["test_param"] = jnp.array([1.0])
    a_constrainers, a_unconstrainers = build_transforms(augmented_params)
    assert "test_param" in list(a_constrainers.keys())
    assert "test_param" in list(a_unconstrainers.keys())
    assert a_constrainers["test_param"](1.0) == 1.0
    assert a_unconstrainers["test_param"](1.0) == 1.0

tests/test_mean_functions.py METASEP
import typing as tp

import jax.numpy as jnp
import pytest
from numpy import isin

from gpjax.mean_functions import Constant, Zero
from gpjax.parameters import initialise


@pytest.mark.parametrize("meanf", [Zero, Constant])
@pytest.mark.parametrize("dim", [1, 2, 5])
def test_shape(meanf, dim):
    meanf = meanf(output_dim=dim)
    x = jnp.linspace(-1.0, 1.0, num=10).reshape(-1, 1)
    if dim > 1:
        x = jnp.hstack([x] * dim)
    params, _, _ = initialise(meanf)
    mu = meanf(x, params)
    assert mu.shape[0] == x.shape[0]
    assert mu.shape[1] == dim


@pytest.mark.parametrize("meanf", [Zero, Constant])
def test_initialisers(meanf):
    params, _, _ = initialise(meanf())
    assert isinstance(params, tp.Dict)

tests/test_likelihoods.py METASEP
import typing as tp

import jax.numpy as jnp
import jax.random as jr
import pytest
import tensorflow_probability.substrates.jax.distributions as tfd

from gpjax.likelihoods import Bernoulli, Gaussian, Likelihood
from gpjax.parameters import initialise

true_initialisation = {
    "Gaussian": ["obs_noise"],
    "Bernoulli": [],
}


@pytest.mark.parametrize("num_datapoints", [1, 10])
@pytest.mark.parametrize("lik", [Gaussian, Bernoulli])
def test_initialisers(num_datapoints, lik):
    l = lik(num_datapoints=num_datapoints)
    params, _, _ = initialise(l)
    assert list(params.keys()) == true_initialisation[l.name]
    assert len(list(params.values())) == len(true_initialisation[l.name])


@pytest.mark.parametrize("n", [1, 10])
def test_predictive_moment(n):
    l = Bernoulli(num_datapoints=n)
    key = jr.PRNGKey(123)
    fmean = jr.uniform(key=key, shape=(n,)) * -1
    fvar = jr.uniform(key=key, shape=(n,))
    pred_mom_fn = l.predictive_moment_fn
    rv = pred_mom_fn(fmean, fvar)
    mu = rv.mean()
    sigma = rv.variance()
    assert isinstance(l.predictive_moment_fn, tp.Callable)
    assert mu.shape == (n,)
    assert sigma.shape == (n,)


@pytest.mark.parametrize("lik", [Gaussian, Bernoulli])
@pytest.mark.parametrize("n", [1, 10])
def test_link_fns(lik: Likelihood, n: int):
    l = lik(num_datapoints=n)
    link_fn = l.link_function
    assert isinstance(link_fn, tp.Callable)
    x = jnp.linspace(-3.0, 3.0).reshape(-1, 1)
    l_eval = link_fn(x)
    if lik == Gaussian:
        assert isinstance(l_eval, jnp.DeviceArray)
    else:
        assert isinstance(l_eval, tfd.Distribution)

tests/test_kernels.py METASEP
from itertools import permutations

import jax.numpy as jnp
import jax.random as jr
import pytest

from gpjax import kernels
from gpjax.kernels import (
    RBF,
    Kernel,
    Matern12,
    Matern32,
    Matern52,
    Polynomial,
    ProductKernel,
    SumKernel,
    cross_covariance,
    euclidean_distance,
    gram,
)
from gpjax.parameters import initialise
from gpjax.utils import I


@pytest.mark.parametrize("kern", [RBF(), Matern12(), Matern32(), Matern52()])
@pytest.mark.parametrize("dim", [1, 2, 5])
def test_gram(kern, dim):
    x = jnp.linspace(-1.0, 1.0, num=10).reshape(-1, 1)
    if dim > 1:
        x = jnp.hstack([x] * dim)
    params, _, _ = initialise(kern)
    gram_matrix = gram(kern, x, params)
    assert gram_matrix.shape[0] == x.shape[0]
    assert gram_matrix.shape[0] == gram_matrix.shape[1]


@pytest.mark.parametrize("kern", [RBF(), Matern12(), Matern32(), Matern52()])
@pytest.mark.parametrize("n1", [3, 10, 20])
@pytest.mark.parametrize("n2", [3, 10, 20])
def test_cross_covariance(kern, n1, n2):
    x1 = jnp.linspace(-1.0, 1.0, num=n1).reshape(-1, 1)
    x2 = jnp.linspace(-1.0, 1.0, num=n2).reshape(-1, 1)
    params, _, _ = initialise(kern)
    kernel_matrix = cross_covariance(kern, x2, x1, params)
    assert kernel_matrix.shape == (n1, n2)


@pytest.mark.parametrize("kernel", [RBF(), Matern12(), Matern32(), Matern52()])
def test_call(kernel):
    params, _, _ = initialise(kernel)
    x, y = jnp.array([[1.0]]), jnp.array([[0.5]])
    point_corr = kernel(x, y, params)
    assert isinstance(point_corr, jnp.DeviceArray)
    assert point_corr.shape == ()


@pytest.mark.parametrize("kern", [RBF(), Matern12(), Matern32(), Matern52()])
@pytest.mark.parametrize("dim", [1, 2, 5])
@pytest.mark.parametrize("ell, sigma", [(0.1, 0.1), (0.5, 0.1), (0.1, 0.5), (0.5, 0.5)])
def test_pos_def(kern, dim, ell, sigma):
    n = 30
    x = jnp.linspace(0.0, 1.0, num=n).reshape(-1, 1)
    if dim > 1:
        x = jnp.hstack((x) * dim)
    params = {"lengthscale": jnp.array([ell]), "variance": jnp.array(sigma)}

    gram_matrix = gram(kern, x, params)
    jitter_matrix = I(n) * 1e-6
    gram_matrix += jitter_matrix
    min_eig = jnp.linalg.eigvals(gram_matrix).min()
    assert min_eig > 0


@pytest.mark.parametrize("kernel", [RBF, Matern12, Matern32, Matern52])
@pytest.mark.parametrize("dim", [1, 2, 5, 10])
def test_initialisation(kernel, dim):
    kern = kernel(active_dims=[i for i in range(dim)])
    params, _, _ = initialise(kern)
    assert list(params.keys()) == ["lengthscale", "variance"]
    assert all(params["lengthscale"] == jnp.array([1.0] * dim))
    assert params["variance"] == jnp.array([1.0])
    if dim > 1:
        assert kern.ard
    else:
        assert not kern.ard


@pytest.mark.parametrize("kernel", [RBF, Matern12, Matern32, Matern52])
def test_dtype(kernel):
    params, _, _ = initialise(kernel())
    for k, v in params.items():
        assert v.dtype == jnp.float64


@pytest.mark.parametrize("degree", [1, 2, 3])
@pytest.mark.parametrize("dim", [1, 2, 5])
@pytest.mark.parametrize("variance", [0.1, 1.0, 2.0])
@pytest.mark.parametrize("shift", [1e-6, 0.1, 1.0])
def test_polynomial(degree, dim, variance, shift):
    x = jnp.linspace(0.0, 1.0, num=20).reshape(-1, 1)
    if dim > 1:
        x = jnp.hstack([x] * dim)
    kern = Polynomial(degree=degree, active_dims=[i for i in range(dim)])
    params = kern.params
    params["shift"] * shift
    params["variance"] * variance
    gram_matrix = gram(kern, x, params)
    assert kern.name == f"Polynomial Degree: {degree}"
    jitter_matrix = I(20) * 1e-6
    gram_matrix += jitter_matrix
    min_eig = jnp.linalg.eigvals(gram_matrix).min()
    assert min_eig > 0
    assert gram_matrix.shape[0] == x.shape[0]
    assert gram_matrix.shape[0] == gram_matrix.shape[1]
    assert list(params.keys()) == ["shift", "variance"]


def test_euclidean_distance():
    x1 = jnp.array(1.0)
    x2 = jnp.array(-4.0)
    x1vec = jnp.array((1, 2, 3))
    x2vec = jnp.array((1, 1, 1))
    assert euclidean_distance(x1vec, x2vec) == 3.0
    assert euclidean_distance(x1, x2) == 5.0


@pytest.mark.parametrize("kernel", [RBF, Matern12, Matern32, Matern52])
def test_active_dim(kernel):
    dim_list = [0, 1, 2, 3]
    perm_length = 2
    dim_pairs = list(permutations(dim_list, r=perm_length))
    n_dims = len(dim_list)
    key = jr.PRNGKey(123)
    X = jr.normal(key, shape=(20, n_dims))

    for dp in dim_pairs:
        Xslice = X[..., dp]
        ad_kern = kernel(active_dims=dp)
        manual_kern = kernel(active_dims=[i for i in range(perm_length)])

        k1 = gram(ad_kern, X, ad_kern.params)
        k2 = gram(manual_kern, Xslice, manual_kern.params)
        assert jnp.all(k1 == k2)


@pytest.mark.parametrize("combination_type", [SumKernel, ProductKernel])
@pytest.mark.parametrize("kernel", [RBF, Matern12, Matern32, Matern52, Polynomial])
@pytest.mark.parametrize("n_kerns", [2, 3, 4])
def test_combination_kernel(combination_type, kernel, n_kerns):
    n = 20
    kern_list = [kernel() for _ in range(n_kerns)]
    c_kernel = combination_type(kernel_set=kern_list)
    assert len(c_kernel.kernel_set) == n_kerns
    assert len(c_kernel.params) == n_kerns
    assert isinstance(c_kernel.kernel_set, list)
    assert isinstance(c_kernel.kernel_set[0], Kernel)
    assert isinstance(c_kernel.params[0], dict)
    x = jnp.linspace(0.0, 1.0, num=n).reshape(-1, 1)
    Kff = gram(c_kernel, x, c_kernel.params)
    assert Kff.shape[0] == Kff.shape[1]
    assert Kff.shape[1] == n


@pytest.mark.parametrize(
    "k1", [RBF(), Matern12(), Matern32(), Matern52(), Polynomial()]
)
@pytest.mark.parametrize(
    "k2", [RBF(), Matern12(), Matern32(), Matern52(), Polynomial()]
)
def test_sum_kern_value(k1, k2):
    n = 10
    sum_kernel = SumKernel(kernel_set=[k1, k2])
    x = jnp.linspace(0.0, 1.0, num=n).reshape(-1, 1)
    Kff = gram(sum_kernel, x, sum_kernel.params)
    Kff_manual = gram(k1, x, k1.params) + gram(k2, x, k2.params)
    assert jnp.all(Kff == Kff_manual)


@pytest.mark.parametrize(
    "k1", [RBF(), Matern12(), Matern32(), Matern52(), Polynomial()]
)
@pytest.mark.parametrize(
    "k2", [RBF(), Matern12(), Matern32(), Matern52(), Polynomial()]
)
def test_prod_kern_value(k1, k2):
    n = 10
    sum_kernel = ProductKernel(kernel_set=[k1, k2])
    x = jnp.linspace(0.0, 1.0, num=n).reshape(-1, 1)
    Kff = gram(sum_kernel, x, sum_kernel.params)
    Kff_manual = gram(k1, x, k1.params) * gram(k2, x, k2.params)
    assert jnp.all(Kff == Kff_manual)


@pytest.mark.parametrize("kernel", [RBF, Matern12, Matern32, Matern52, Polynomial])
def test_update_params(kernel):
    kern = kernel()
    params = kern.params
    for k in params.keys():
        params[k] = 100
        kern._params = params
        assert kern.params == params

tests/test_gp.py METASEP
import typing as tp
from turtle import pos

import jax.numpy as jnp
import jax.random as jr
import pytest
import tensorflow_probability.substrates.jax as tfp
from numpy import isin

from gpjax import Dataset, initialise, likelihoods, transform
from gpjax.gps import (
    GP,
    ConjugatePosterior,
    NonConjugatePosterior,
    Prior,
    construct_posterior,
)
from gpjax.kernels import RBF
from gpjax.likelihoods import Bernoulli, Gaussian, NonConjugateLikelihoods
from gpjax.parameters import initialise


@pytest.mark.parametrize("num_datapoints", [1, 10])
def test_prior(num_datapoints):
    p = Prior(kernel=RBF())
    params, _, _ = initialise(p)
    assert isinstance(p, Prior)
    assert isinstance(p, GP)
    meanf = p.mean(params)
    varf = p.variance(params)
    assert isinstance(meanf, tp.Callable)
    assert isinstance(varf, tp.Callable)
    x = jnp.linspace(-3.0, 3.0, num_datapoints).reshape(-1, 1)
    mu = meanf(x)
    sigma = varf(x)
    assert mu.shape == (num_datapoints, 1)
    assert sigma.shape == (num_datapoints, num_datapoints)

    rv = p.random_variable(x, params)
    assert isinstance(rv, tfp.distributions.Distribution)


@pytest.mark.parametrize("num_datapoints", [2, 10])
def test_conjugate_posterior(num_datapoints):
    key = jr.PRNGKey(123)
    x = jnp.sort(
        jr.uniform(key=key, minval=-2.0, maxval=2.0, shape=(num_datapoints, 1)),
        axis=0,
    )
    y = jnp.sin(x) + jr.normal(key=key, shape=x.shape) * 0.1
    D = Dataset(X=x, y=y)
    # Initialisation
    p = Prior(kernel=RBF())
    lik = Gaussian(num_datapoints=num_datapoints)
    post = p * lik
    assert isinstance(post, ConjugatePosterior)
    assert isinstance(post, GP)
    assert isinstance(p, GP)
    params, constrainers, unconstrainer = initialise(post)
    params = transform(params, unconstrainer)

    # Marginal likelihood
    mll = post.marginal_log_likelihood(training=D, transformations=constrainers)
    objective_val = mll(params)
    assert isinstance(objective_val, jnp.DeviceArray)
    assert objective_val.shape == ()

    # Prediction
    meanf = post.mean(training_data=D, params=params)
    varf = post.variance(training_data=D, params=params)
    assert isinstance(meanf, tp.Callable)
    assert isinstance(varf, tp.Callable)
    x = jnp.linspace(-3.0, 3.0, num_datapoints).reshape(-1, 1)
    mu = meanf(x)
    sigma = varf(x)
    assert mu.shape == (num_datapoints, 1)
    assert sigma.shape == (num_datapoints, num_datapoints)


@pytest.mark.parametrize("num_datapoints", [2, 10])
@pytest.mark.parametrize("likel", NonConjugateLikelihoods)
def test_nonconjugate_posterior(num_datapoints, likel):
    key = jr.PRNGKey(123)
    x = jnp.sort(
        jr.uniform(key=key, minval=-2.0, maxval=2.0, shape=(num_datapoints, 1)),
        axis=0,
    )
    y = 0.5 * jnp.sign(jnp.cos(3 * x + jr.normal(key, shape=x.shape) * 0.05)) + 0.5
    D = Dataset(X=x, y=y)
    # Initialisation
    p = Prior(kernel=RBF())
    lik = likel(num_datapoints=num_datapoints)
    post = p * lik
    assert isinstance(post, NonConjugatePosterior)
    assert isinstance(post, GP)
    assert isinstance(p, GP)
    params, constrainers, unconstrainer = initialise(post)
    params = transform(params, unconstrainer)

    # Marginal likelihood
    mll = post.marginal_log_likelihood(training=D, transformations=constrainers)
    objective_val = mll(params)
    assert isinstance(objective_val, jnp.DeviceArray)
    assert objective_val.shape == ()

    # Prediction
    meanf = post.mean(training_data=D, params=params)
    varf = post.variance(training_data=D, params=params)
    assert isinstance(meanf, tp.Callable)
    assert isinstance(varf, tp.Callable)
    x = jnp.linspace(-3.0, 3.0, num_datapoints).reshape(-1, 1)
    mu = meanf(x)
    sigma = varf(x)
    assert mu.shape == (num_datapoints, 1)
    assert sigma.shape == (num_datapoints, num_datapoints)


@pytest.mark.parametrize("num_datapoints", [1, 10])
@pytest.mark.parametrize("lik", [Bernoulli, Gaussian])
def test_param_construction(num_datapoints, lik):
    p = Prior(kernel=RBF()) * lik(num_datapoints=num_datapoints)
    params, _, _ = initialise(p)
    if isinstance(lik, Bernoulli):
        assert sorted(list(params.keys())) == [
            "kernel",
            "latent_fn",
            "likelihood",
            "mean_function",
        ]
    elif isinstance(lik, Gaussian):
        assert sorted(list(params.keys())) == [
            "kernel",
            "likelihood",
            "mean_function",
        ]


@pytest.mark.parametrize("lik", [Bernoulli, Gaussian])
def test_posterior_construct(lik):
    pr = Prior(kernel=RBF())
    l = lik(num_datapoints=10)
    p1 = pr * l
    p2 = construct_posterior(prior=pr, likelihood=l)
    assert type(p1) == type(p2)

tests/test_config.py METASEP
from ml_collections import ConfigDict
from tensorflow_probability.substrates.jax import bijectors as tfb

from gpjax.config import add_parameter, get_defaults


def test_add_parameter():
    add_parameter("test_parameter", tfb.Identity())
    config = get_defaults()
    assert "test_parameter" in config.transformations
    assert "test_parameter_transform" in config.transformations
    assert config.transformations["test_parameter"] == "test_parameter_transform"
    assert isinstance(config.transformations["test_parameter_transform"], tfb.Bijector)


def test_add_parameter():
    config = get_defaults()
    add_parameter("test_parameter", tfb.Identity())
    config = get_defaults()
    assert "test_parameter" in config.transformations
    assert "test_parameter_transform" in config.transformations
    assert config.transformations["test_parameter"] == "test_parameter_transform"
    assert isinstance(config.transformations["test_parameter_transform"], tfb.Bijector)


def test_get_defaults():
    config = get_defaults()
    assert isinstance(config, ConfigDict)
    assert isinstance(config.transformations, ConfigDict)

tests/test_abstractions.py METASEP
import jax.numpy as jnp
import jax.random as jr
import optax
import pytest
from jax.experimental import optimizers

from gpjax import RBF, Dataset, Gaussian, Prior, initialise, transform
from gpjax.abstractions import fit, optax_fit


@pytest.mark.parametrize("n", [20])
def test_fit(n):
    key = jr.PRNGKey(123)
    x = jnp.sort(jr.uniform(key=key, minval=-2.0, maxval=2.0, shape=(n, 1)), axis=0)
    y = jnp.sin(x) + jr.normal(key=key, shape=x.shape) * 0.1
    D = Dataset(X=x, y=y)
    p = Prior(kernel=RBF()) * Gaussian(num_datapoints=n)
    params, constrainer, unconstrainer = initialise(p)
    mll = p.marginal_log_likelihood(D, constrainer, negative=True)
    pre_mll_val = mll(params)
    opt_init, opt_update, get_params = optimizers.adam(step_size=0.1)
    optimised_params = fit(mll, params, opt_init, opt_update, get_params, n_iters=10)
    optimised_params = transform(optimised_params, constrainer)
    assert isinstance(optimised_params, dict)
    assert mll(optimised_params) < pre_mll_val


@pytest.mark.parametrize("n", [20])
def test_optax_fit(n):
    key = jr.PRNGKey(123)
    x = jnp.sort(jr.uniform(key=key, minval=-2.0, maxval=2.0, shape=(n, 1)), axis=0)
    y = jnp.sin(x) + jr.normal(key=key, shape=x.shape) * 0.1
    D = Dataset(X=x, y=y)
    p = Prior(kernel=RBF()) * Gaussian(num_datapoints=n)
    params, constrainer, unconstrainer = initialise(p)
    mll = p.marginal_log_likelihood(D, constrainer, negative=True)
    pre_mll_val = mll(params)
    optimiser = optax.adam(learning_rate=0.1)
    optimised_params = optax_fit(mll, params, optimiser, n_iters=10)
    optimised_params = transform(optimised_params, constrainer)
    assert isinstance(optimised_params, dict)
    assert mll(optimised_params) < pre_mll_val

tests/__init__.py METASEP

gpjax/utils.py METASEP
import typing as tp
from copy import deepcopy

import jax
import jax.numpy as jnp
from jax.scipy.linalg import cho_factor, cho_solve

from .types import Array


def I(n: int) -> Array:
    """
    Compute an n x n identity matrix.
    :param n: The size of of the matrix.
    :return: An n x n identity matrix.
    """
    return jnp.eye(n)


def concat_dictionaries(a: dict, b: dict) -> dict:
    """
    Append one dictionary below another. If duplicate keys exist, then the key-value pair of the second supplied
    dictionary will be used.
    """
    return {**a, **b}


def merge_dictionaries(base_dict: dict, in_dict: dict) -> dict:
    """
    This will return a complete dictionary based on the keys of the first matrix. If the same key should exist in the
    second matrix, then the key-value pair from the first dictionary will be overwritten. The purpose of this is that
    the base_dict will be a complete dictionary of values such that an incomplete second dictionary can be used to
    update specific key-value pairs.

    :param base_dict: Complete dictionary of key-value pairs.
    :param in_dict: Subset of key-values pairs such that values from this dictionary will take precedent.
    :return: A merged single dictionary.
    """
    for k, v in base_dict.items():
        if k in in_dict.keys():
            base_dict[k] = in_dict[k]
    return base_dict


def sort_dictionary(base_dict: dict) -> dict:
    """
    Sort a dictionary based on the dictionary's key values.

    :param base_dict: The unsorted dictionary.
    :return: A dictionary sorted alphabetically on the dictionary's keys.
    """
    return dict(sorted(base_dict.items()))


def as_constant(parameter_set: dict, params: list) -> tp.Tuple[dict, dict]:
    base_params = deepcopy(parameter_set)
    sparams = {}
    for param in params:
        sparams[param] = base_params[param]
        del base_params[param]
    return base_params, sparams


def dict_array_coercion(params) -> tp.Tuple[tp.Callable, tp.Callable]:
    flattened_pytree = jax.tree_util.tree_flatten(params)

    def dict_to_array(parameter_dict) -> jnp.DeviceArray:
        return jax.tree_util.tree_flatten(parameter_dict)[0]

    def array_to_dict(parameter_array) -> tp.Dict:
        return jax.tree_util.tree_unflatten(flattened_pytree[1], parameter_array)

    return dict_to_array, array_to_dict

gpjax/types.py METASEP
import typing as tp

import jax.numpy as jnp
import numpy as np
from chex import dataclass

NoneType = type(None)
Array = tp.Union[np.ndarray, jnp.ndarray]


@dataclass(repr=False)
class Dataset:
    X: Array
    y: Array = None

    def __repr__(self) -> str:
        return (
            f"- Number of datapoints: {self.X.shape[0]}\n- Dimension:"
            f" {self.X.shape[1]}"
        )

    @property
    def n(self) -> int:
        return self.X.shape[0]

    @property
    def in_dim(self) -> int:
        return self.X.shape[1]

    @property
    def out_dim(self) -> int:
        return self.y.shape[1]


def verify_dataset(ds: Dataset) -> None:
    assert ds.X.ndim == 2, (
        "2-dimensional training inputs are required. Current dimension:"
        f" {ds.X.ndim}."
    )
    if ds.y is not None:
        assert ds.y.ndim == 2, (
            "2-dimensional training outputs are required. Current dimension:"
            f" {ds.y.ndim}."
        )
        assert ds.X.shape[0] == ds.y.shape[0], (
            "Number of inputs must equal the number of outputs. \nCurrent"
            f" counts:\n- X: {ds.X.shape[0]}\n- y: {ds.y.shape[0]}"
        )

gpjax/parameters.py METASEP
import typing as tp
import warnings
from copy import deepcopy

import jax
import jax.numpy as jnp
import tensorflow_probability.substrates.jax as tfp
from tensorflow_probability.substrates.jax import distributions as tfd

from .config import get_defaults
from .types import Array


################################
# Base operations
################################
def initialise(obj) -> tp.Tuple[tp.Dict, tp.Dict, tp.Dict]:
    params = obj.params
    constrainers, unconstrainers = build_transforms(params)
    return params, constrainers, unconstrainers


def recursive_items(d1, d2):
    for key, value in d1.items():
        if type(value) is dict:
            yield from recursive_items(value, d2[key])
        else:
            yield (key, value, d2[key])


def recursive_complete(d1, d2) -> dict:
    for key, value in d1.items():
        if type(value) is dict:
            if key in d2.keys():
                recursive_complete(value, d2[key])
            # else:
            #     pass
        else:
            if key in d2.keys():
                d1[key] = d2[key]
    return d1


# def recursive_fn(d1, d2, fn: tp.Callable[[tp.Any], tp.Any]):
#     for key, value in d1.items():
#         if type(value) is dict:
#             yield from recursive_fn(value, d2[key], fn)
#         else:
#             yield fn(value, d2[key])


################################
# Parameter transformation
################################
def build_transforms(params, key=None) -> tp.Tuple[tp.Dict, tp.Dict]:
    constrainer, unconstrainer = copy_dict_structure(params), copy_dict_structure(
        params
    )
    config = get_defaults()
    transform_set = config["transformations"]

    def recursive_transforms(ps, cs, ucs) -> tp.Tuple[tp.Dict, tp.Dict]:
        for key, value in ps.items():
            if type(value) is dict:
                recursive_transforms(value, cs[key], ucs[key])
            else:
                if key in transform_set.keys():
                    transform_type = transform_set[key]
                    bijector = transform_set[transform_type]
                else:
                    bijector = tfp.bijectors.Identity()
                    warnings.warn(
                        f"Parameter {key} has no transform. Defaulting to"
                        " identity transfom."
                    )
                cs[key] = bijector.forward
                ucs[key] = bijector.inverse
        return cs, ucs

    constrainers, unconstrainers = recursive_transforms(
        params, constrainer, unconstrainer
    )
    return constrainers, unconstrainers


def transform(params: dict, transform_map: dict):
    transformed_params = copy_dict_structure(params)

    def apply_transform(untransformed_params, transformed_params, transform_map):
        for key, value in untransformed_params.items():
            if type(value) is dict:
                apply_transform(value, transformed_params[key], transform_map[key])
            else:
                transformed_params[key] = transform_map[key](value)
        return transformed_params

    return apply_transform(params, transformed_params, transform_map)


################################
# Priors
################################
def log_density(param: jnp.DeviceArray, density: tfd.Distribution) -> Array:
    if type(density) == type(None):
        log_prob = jnp.array(0.0)
    else:
        log_prob = jnp.sum(density.log_prob(param))
    return log_prob


def copy_dict_structure(params: dict) -> dict:
    # Copy dictionary structure
    prior_container = deepcopy(params)
    # Set all values to zero
    prior_container = jax.tree_map(lambda _: None, prior_container)
    return prior_container


def structure_priors(params: dict, priors: dict) -> dict:
    """First create a dictionary with equal structure to the parameters. Then, for each supplied prior, overwrite the None value if it exists.

    Args:
        params (dict): [description]
        priors (dict): [description]

    Returns:
        dict: [description]
    """
    prior_container = copy_dict_structure(params)
    # Where a prior has been supplied, override the None value by the prior distribution.
    complete_prior = recursive_complete(prior_container, priors)
    return complete_prior


def evaluate_priors(params: dict, priors: dict) -> dict:
    """Recursive loop over pair of dictionaries that correspond to a parameter's
    current value and the parameter's respective prior distribution. For
    parameters where a prior distribution is specified, the log-prior density is
    evaluated at the parameter's current value.

    Args: params (dict): Dictionary containing the current set of parameter
        estimates. priors (dict): Dictionary specifying the parameters' prior
        distributions.

    Returns: Array: The log-prior density, summed over all parameters.
    """
    lpd = jnp.array(0.0)
    if priors is not None:
        for name, param, prior in recursive_items(params, priors):
            lpd += log_density(param, prior)
    return lpd


def prior_checks(priors: dict) -> dict:
    if "latent" in priors.keys():
        latent_prior = priors["latent"]
        if isinstance(latent_prior, tfd.Distribution) and latent_prior.name != "Normal":
            warnings.warn(
                f"A {latent_prior.name} distribution prior has been placed on"
                " the latent function. It is strongly advised that a"
                " unit-Gaussian prior is used."
            )
        else:
            if not latent_prior:
                priors["latent"] = tfd.Normal(loc=0.0, scale=1.0)
    else:
        priors["latent"] = tfd.Normal(loc=0.0, scale=1.0)

    return priors

gpjax/mean_functions.py METASEP
import abc
from typing import Dict, Optional

import jax.numpy as jnp
from chex import dataclass

from .types import Array


@dataclass(repr=False)
class MeanFunction:
    output_dim: Optional[int] = 1
    name: Optional[str] = "Mean function"

    @abc.abstractmethod
    def __call__(self, x: Array) -> Array:
        raise NotImplementedError

    def __repr__(self):
        return f"{self.name}\n\t Output dimension: {self.output_dim}"

    @property
    @abc.abstractmethod
    def params(self) -> dict:
        raise NotImplementedError


@dataclass(repr=False)
class Zero(MeanFunction):
    output_dim: Optional[int] = 1
    name: Optional[str] = "Zero mean function"

    def __call__(self, x: Array, params: dict) -> Array:
        out_shape = (x.shape[0], self.output_dim)
        return jnp.zeros(shape=out_shape)

    @property
    def params(self) -> dict:
        return {}


@dataclass(repr=False)
class Constant(MeanFunction):
    output_dim: Optional[int] = 1
    name: Optional[str] = "Constant mean function"

    def __call__(self, x: Array, params: Dict) -> Array:
        out_shape = (x.shape[0], self.output_dim)
        return jnp.ones(shape=out_shape) * params["variance"]

    @property
    def params(self) -> dict:
        return {"variance": jnp.array(1.0)}

gpjax/likelihoods.py METASEP
import abc
from typing import Callable, Optional

import jax.numpy as jnp
from chex import dataclass
from tensorflow_probability.substrates.jax import distributions as tfd


@dataclass(repr=False)
class Likelihood:
    num_datapoints: int  # The number of datapoints that the likelihood factorises over
    name: Optional[str] = "Likelihood"

    def __repr__(self):
        return f"{self.name} likelihood function"

    @property
    @abc.abstractmethod
    def params(self) -> dict:
        raise NotImplementedError

    @property
    @abc.abstractmethod
    def link_function(self) -> Callable:
        raise NotImplementedError


@dataclass(repr=False)
class Gaussian(Likelihood):
    name: Optional[str] = "Gaussian"

    @property
    def params(self) -> dict:
        return {"obs_noise": jnp.array([1.0])}

    @property
    def link_function(self) -> Callable:
        def identity_fn(x):
            return x

        return identity_fn


@dataclass(repr=False)
class Bernoulli(Likelihood):
    name: Optional[str] = "Bernoulli"

    @property
    def params(self) -> dict:
        return {}

    @property
    def link_function(self) -> Callable:
        def link_fn(x):
            return tfd.ProbitBernoulli(x)

        return link_fn

    @property
    def predictive_moment_fn(self) -> Callable:
        def moment_fn(mean: jnp.DeviceArray, variance: jnp.DeviceArray):
            rv = self.link_function(mean / jnp.sqrt(1 + variance))
            return rv

        return moment_fn


NonConjugateLikelihoods = [Bernoulli]
NonConjugateLikelihoodType = Bernoulli  # Union[Bernoulli]

gpjax/kernels.py METASEP
import abc
from re import L
from typing import Callable, Dict, List, Optional, Tuple
from unicodedata import name

import jax.numpy as jnp
from chex import dataclass
from jax import vmap

from gpjax.types import Array


@dataclass(repr=False)
class Kernel:
    active_dims: Optional[List[int]] = None
    stationary: Optional[bool] = False
    spectral: Optional[bool] = False
    name: Optional[str] = "Kernel"
    _params: Optional[Dict] = None

    def __post_init__(self):
        self.ndims = 1 if not self.active_dims else len(self.active_dims)

    @abc.abstractmethod
    def __call__(self, x: Array, y: Array, params: dict) -> Array:
        raise NotImplementedError

    def slice_input(self, x: Array) -> Array:
        return x[..., self.active_dims]

    @property
    def ard(self):
        return True if self.ndims > 1 else False

    @property
    def params(self) -> dict:
        return self._params

    @params.setter
    def params(self, value):
        self._params = value


@dataclass
class CombinationKernel:
    kernel_set: List[Kernel]
    name: Optional[str] = "Combination kernel"
    combination_fn: Optional[Callable] = None

    @property
    def params(self) -> List[Dict]:
        return [kernel.params for kernel in self.kernel_set]

    def __call__(self, x: Array, y: Array, params: dict) -> Array:
        return self.combination_fn(
            jnp.stack([k(x, y, p) for k, p in zip(self.kernel_set, params)])
        )


@dataclass
class SumKernel(CombinationKernel):
    name: Optional[str] = "Sum kernel"
    combination_fn: Optional[Callable] = jnp.sum


@dataclass
class ProductKernel(CombinationKernel):
    name: Optional[str] = "Product kernel"
    combination_fn: Optional[Callable] = jnp.prod


@dataclass(repr=False)
class RBF(Kernel):
    name: Optional[str] = "Radial basis function kernel"

    def __post_init__(self):
        self.ndims = 1 if not self.active_dims else len(self.active_dims)
        self._params = {
            "lengthscale": jnp.array([1.0] * self.ndims),
            "variance": jnp.array([1.0]),
        }

    def __call__(self, x: jnp.DeviceArray, y: jnp.DeviceArray, params: dict) -> Array:
        x = self.slice_input(x) / params["lengthscale"]
        y = self.slice_input(y) / params["lengthscale"]
        K = params["variance"] * jnp.exp(-0.5 * squared_distance(x, y))
        return K.squeeze()


@dataclass(repr=False)
class Matern12(Kernel):
    name: Optional[str] = "Matern 1/2"

    def __post_init__(self):
        self.ndims = 1 if not self.active_dims else len(self.active_dims)
        self._params = {
            "lengthscale": jnp.array([1.0] * self.ndims),
            "variance": jnp.array([1.0]),
        }

    def __call__(self, x: jnp.DeviceArray, y: jnp.DeviceArray, params: dict) -> Array:
        x = self.slice_input(x) / params["lengthscale"]
        y = self.slice_input(y) / params["lengthscale"]
        K = params["variance"] * jnp.exp(-0.5 * euclidean_distance(x, y))
        return K.squeeze()


@dataclass(repr=False)
class Matern32(Kernel):
    name: Optional[str] = "Matern 3/2"

    def __post_init__(self):
        self.ndims = 1 if not self.active_dims else len(self.active_dims)
        self._params = {
            "lengthscale": jnp.array([1.0] * self.ndims),
            "variance": jnp.array([1.0]),
        }

    def __call__(self, x: jnp.DeviceArray, y: jnp.DeviceArray, params: dict) -> Array:
        x = self.slice_input(x) / params["lengthscale"]
        y = self.slice_input(y) / params["lengthscale"]
        tau = euclidean_distance(x, y)
        K = (
            params["variance"]
            * (1.0 + jnp.sqrt(3.0) * tau)
            * jnp.exp(-jnp.sqrt(3.0) * tau)
        )
        return K.squeeze()


@dataclass(repr=False)
class Matern52(Kernel):
    name: Optional[str] = "Matern 5/2"

    def __post_init__(self):
        self.ndims = 1 if not self.active_dims else len(self.active_dims)
        self._params = {
            "lengthscale": jnp.array([1.0] * self.ndims),
            "variance": jnp.array([1.0]),
        }

    def __call__(self, x: jnp.DeviceArray, y: jnp.DeviceArray, params: dict) -> Array:
        x = self.slice_input(x) / params["lengthscale"]
        y = self.slice_input(y) / params["lengthscale"]
        tau = euclidean_distance(x, y)
        K = (
            params["variance"]
            * (1.0 + jnp.sqrt(5.0) * tau + 5.0 / 3.0 * jnp.square(tau))
            * jnp.exp(-jnp.sqrt(5.0) * tau)
        )
        return K.squeeze()


@dataclass(repr=False)
class Polynomial(Kernel):
    name: Optional[str] = "Polynomial"
    degree: int = 1

    def __post_init__(self):
        self.ndims = 1 if not self.active_dims else len(self.active_dims)
        self._params = {
            "shift": jnp.array([1.0]),
            "variance": jnp.array([1.0] * self.ndims),
        }
        self.name = f"Polynomial Degree: {self.degree}"

    def __call__(self, x: jnp.DeviceArray, y: jnp.DeviceArray, params: dict) -> Array:
        x = self.slice_input(x).squeeze()
        y = self.slice_input(y).squeeze()
        K = jnp.power(params["shift"] + jnp.dot(x * params["variance"], y), self.degree)
        return K.squeeze()


def squared_distance(x: Array, y: Array):
    return jnp.sum((x - y) ** 2)


def euclidean_distance(x: Array, y: Array):
    return jnp.sum(jnp.abs(x - y))


def gram(kernel: Kernel, inputs: Array, params: dict) -> Array:
    return vmap(lambda x1: vmap(lambda y1: kernel(x1, y1, params))(inputs))(inputs)


def cross_covariance(kernel: Kernel, x: Array, y: Array, params: dict) -> Array:
    return vmap(lambda x1: vmap(lambda y1: kernel(x1, y1, params))(x))(y)

gpjax/gps.py METASEP
import typing as tp
from abc import abstractmethod, abstractproperty

import jax.numpy as jnp
import tensorflow_probability.substrates.jax.distributions as tfd
from chex import dataclass
from jax.scipy.linalg import cho_factor, cho_solve, solve_triangular

from gpjax.config import get_defaults

from .kernels import Kernel, cross_covariance, gram
from .likelihoods import (
    Gaussian,
    Likelihood,
    NonConjugateLikelihoods,
    NonConjugateLikelihoodType,
)
from .mean_functions import MeanFunction, Zero
from .parameters import copy_dict_structure, evaluate_priors, transform
from .types import Array, Dataset
from .utils import I, concat_dictionaries


@dataclass
class GP:
    @abstractmethod
    def mean(self) -> tp.Callable[[Dataset], Array]:
        raise NotImplementedError

    @abstractmethod
    def variance(self) -> tp.Callable[[Dataset], Array]:
        raise NotImplementedError

    @abstractproperty
    def params(self) -> tp.Dict:
        raise NotImplementedError


#######################
# GP Priors
#######################
@dataclass(repr=False)
class Prior(GP):
    kernel: Kernel
    mean_function: tp.Optional[MeanFunction] = Zero()
    name: tp.Optional[str] = "Prior"

    def __mul__(self, other: Gaussian):
        return construct_posterior(prior=self, likelihood=other)

    def mean(self, params: dict) -> tp.Callable[[Array], Array]:
        def mean_fn(test_points: Array):
            mu = self.mean_function(test_points, params["mean_function"])
            return mu

        return mean_fn

    def variance(self, params: dict) -> tp.Callable[[Array], Array]:
        def variance_fn(test_points: Array):
            Kff = gram(self.kernel, test_points, params["kernel"])
            jitter_matrix = I(test_points.shape[0]) * 1e-8
            covariance_matrix = Kff + jitter_matrix
            return covariance_matrix

        return variance_fn

    @property
    def params(self) -> dict:
        return {
            "kernel": self.kernel.params,
            "mean_function": self.mean_function.params,
        }

    def random_variable(self, test_points: Array, params: dict) -> tfd.Distribution:
        n = test_points.shape[0]
        mu = self.mean(params)(test_points)
        sigma = self.variance(params)(test_points)
        sigma += I(n) * 1e-8
        return tfd.MultivariateNormalTriL(mu.squeeze(), jnp.linalg.cholesky(sigma))


#######################
# GP Posteriors
#######################
@dataclass
class Posterior(GP):
    prior: Prior
    likelihood: Likelihood
    name: tp.Optional[str] = "GP Posterior"

    @abstractmethod
    def mean(
        self, training_data: Dataset, params: dict
    ) -> tp.Callable[[Dataset], Array]:
        raise NotImplementedError

    @abstractmethod
    def variance(
        self, training_data: Dataset, params: dict
    ) -> tp.Callable[[Dataset], Array]:
        raise NotImplementedError

    @property
    def params(self) -> dict:
        return concat_dictionaries(
            self.prior.params, {"likelihood": self.likelihood.params}
        )


@dataclass
class ConjugatePosterior(Posterior):
    prior: Prior
    likelihood: Gaussian
    name: tp.Optional[str] = "ConjugatePosterior"

    def mean(self, training_data: Dataset, params: dict) -> tp.Callable[[Array], Array]:
        X, y = training_data.X, training_data.y
        sigma = params["likelihood"]["obs_noise"]
        n_train = training_data.n
        # Precompute covariance matrices
        Kff = gram(self.prior.kernel, X, params["kernel"])
        prior_mean = self.prior.mean_function(X, params["mean_function"])
        L = cho_factor(Kff + I(n_train) * sigma, lower=True)

        prior_distance = y - prior_mean
        weights = cho_solve(L, prior_distance)

        def mean_fn(test_inputs: Array) -> Array:
            prior_mean_at_test_inputs = self.prior.mean_function(
                test_inputs, params["mean_function"]
            )
            Kfx = cross_covariance(self.prior.kernel, X, test_inputs, params["kernel"])
            return prior_mean_at_test_inputs + jnp.dot(Kfx, weights)

        return mean_fn

    def variance(
        self, training_data: Dataset, params: dict
    ) -> tp.Callable[[Array], Array]:
        X = training_data.X
        n_train = training_data.n
        obs_noise = params["likelihood"]["obs_noise"]
        n_train = training_data.n
        Kff = gram(self.prior.kernel, X, params["kernel"])
        Kff += I(n_train) * 1e-8
        L = cho_factor(Kff + I(n_train) * obs_noise, lower=True)

        def variance_fn(test_inputs: Array) -> Array:
            Kfx = cross_covariance(self.prior.kernel, X, test_inputs, params["kernel"])
            Kxx = gram(self.prior.kernel, test_inputs, params["kernel"])
            latent_values = cho_solve(L, Kfx.T)
            return Kxx - jnp.dot(Kfx, latent_values)

        return variance_fn

    def marginal_log_likelihood(
        self,
        training: Dataset,
        transformations: tp.Dict,
        priors: dict = None,
        static_params: dict = None,
        negative: bool = False,
    ) -> tp.Callable[[Dataset], Array]:
        x, y = training.X, training.y

        def mll(
            params: dict,
        ):
            params = transform(params=params, transform_map=transformations)
            if static_params:
                params = concat_dictionaries(params, transform(static_params))
            mu = self.prior.mean_function(x, params)
            gram_matrix = gram(self.prior.kernel, x, params["kernel"])
            gram_matrix += params["likelihood"]["obs_noise"] * I(x.shape[0])
            L = jnp.linalg.cholesky(gram_matrix)
            random_variable = tfd.MultivariateNormalTriL(mu, L)

            log_prior_density = evaluate_priors(params, priors)
            constant = jnp.array(-1.0) if negative else jnp.array(1.0)
            return constant * (
                random_variable.log_prob(y.squeeze()).mean() + log_prior_density
            )

        return mll


@dataclass
class NonConjugatePosterior(Posterior):
    prior: Prior
    likelihood: NonConjugateLikelihoodType
    name: tp.Optional[str] = "Non-Conjugate Posterior"

    def __repr__(self):
        mean_fn_string = self.prior.mean_function.__repr__()
        kernel_string = self.prior.kernel.__repr__()
        likelihood_string = self.likelihood.__repr__()
        return (
            f"Conjugate Posterior\n{'-'*80}\n- {mean_fn_string}\n-"
            f" {kernel_string}\n- {likelihood_string}"
        )

    @property
    def params(self) -> dict:
        hyperparameters = concat_dictionaries(
            self.prior.params, {"likelihood": self.likelihood.params}
        )
        hyperparameters["latent"] = jnp.zeros(shape=(self.likelihood.num_datapoints, 1))
        return hyperparameters

    def mean(
        self, training_data: Dataset, params: dict
    ) -> tp.Callable[[Dataset], Array]:
        X, y = training_data.X, training_data.y
        N = training_data.n
        Kff = gram(self.prior.kernel, X, params["kernel"])
        L = jnp.linalg.cholesky(Kff + I(N) * 1e-6)

        def meanf(test_inputs: Array) -> Array:
            Kfx = cross_covariance(self.prior.kernel, X, test_inputs, params["kernel"])
            Kxx = gram(self.prior.kernel, test_inputs, params["kernel"])
            A = solve_triangular(L, Kfx.T, lower=True)
            latent_var = Kxx - jnp.sum(jnp.square(A), -2)
            latent_mean = jnp.matmul(A.T, params["latent"])

            lvar = jnp.diag(latent_var)

            moment_fn = self.likelihood.predictive_moment_fn
            pred_rv = moment_fn(latent_mean.ravel(), lvar)
            return pred_rv.mean().reshape(-1, 1)

        return meanf

    def variance(
        self, training_data: Dataset, params: dict
    ) -> tp.Callable[[Dataset], Array]:
        X, y = training_data.X, training_data.y
        N = training_data.n
        Kff = gram(self.prior.kernel, X, params["kernel"])
        L = jnp.linalg.cholesky(Kff + I(N) * 1e-6)

        def variancef(test_inputs: Array) -> Array:
            Kfx = cross_covariance(self.prior.kernel, X, test_inputs, params["kernel"])
            Kxx = gram(self.prior.kernel, test_inputs, params["kernel"])
            A = solve_triangular(L, Kfx.T, lower=True)
            latent_var = Kxx - jnp.sum(jnp.square(A), -2)
            latent_mean = jnp.matmul(A.T, params["latent"])

            lvar = jnp.diag(latent_var)

            moment_fn = self.likelihood.predictive_moment_fn
            pred_rv = moment_fn(latent_mean.ravel(), lvar)
            return jnp.diag(pred_rv.variance())

        return variancef

    def marginal_log_likelihood(
        self,
        training: Dataset,
        transformations: tp.Dict,
        priors: dict = None,
        static_params: dict = None,
        negative: bool = False,
    ) -> tp.Callable[[Dataset], Array]:
        x, y = training.X, training.y
        n = training.n
        jitter = get_defaults()["jitter"]
        if not priors:
            priors = copy_dict_structure(self.params)
            priors["latent"] = tfd.Normal(loc=0.0, scale=1.0)

        def mll(params: dict):
            params = transform(params=params, transform_map=transformations)
            if static_params:
                params = concat_dictionaries(params, transform(static_params))
            gram_matrix = gram(self.prior.kernel, x, params["kernel"])
            gram_matrix += I(n) * jitter
            L = jnp.linalg.cholesky(gram_matrix)
            F = jnp.matmul(L, params["latent"])
            rv = self.likelihood.link_function(F)
            ll = jnp.sum(rv.log_prob(y))

            log_prior_density = evaluate_priors(params, priors)
            constant = jnp.array(-1.0) if negative else jnp.array(1.0)
            return constant * (ll + log_prior_density)

        return mll


def construct_posterior(prior: Prior, likelihood: Likelihood) -> Posterior:
    if isinstance(likelihood, Gaussian):
        PosteriorGP = ConjugatePosterior
    elif any([isinstance(likelihood, l) for l in NonConjugateLikelihoods]):
        PosteriorGP = NonConjugatePosterior
    else:
        raise NotImplementedError(
            f"No posterior implemented for {likelihood.name} likelihood"
        )
    return PosteriorGP(prior=prior, likelihood=likelihood)

gpjax/config.py METASEP
from typing import Tuple

import jax.random as jr
import tensorflow_probability.substrates.jax.bijectors as tfb
from ml_collections import ConfigDict

__config = None


def get_defaults() -> ConfigDict:
    config = ConfigDict()
    config.key = jr.PRNGKey(123)
    # Covariance matrix stabilising jitter
    config.jitter = 1e-6

    # Default bijections
    config.transformations = transformations = ConfigDict()
    transformations.positive_transform = tfb.Softplus()
    transformations.identity_transform = tfb.Identity()

    # Default parameter transforms
    transformations.lengthscale = "positive_transform"
    transformations.shift = "positive_transform"
    transformations.variance = "positive_transform"
    transformations.obs_noise = "positive_transform"
    transformations.latent = "identity_transform"
    transformations.basis_fns = "identity_transform"
    transformations.offset = "identity_transform"
    global __config
    if not __config:
        __config = config
    return __config


def add_parameter(param_name: str, bijection: tfb.Bijector) -> None:
    lookup_name = f"{param_name}_transform"
    get_defaults()
    __config.transformations[lookup_name] = bijection
    __config.transformations[param_name] = lookup_name

gpjax/abstractions.py METASEP
import typing as tp

import jax
import jax.numpy as jnp
import optax
from tqdm import trange


def fit(
    objective,
    params: dict,
    opt_init,
    opt_update,
    get_params,
    n_iters: int = 100,
    log_rate: int = 10,
) -> tp.Dict:
    opt_state = opt_init(params)

    def step(i, opt_state):
        params = get_params(opt_state)
        v, g = jax.value_and_grad(objective)(params)
        return opt_update(i, g, opt_state), v

    tr = trange(n_iters)
    for i in tr:
        opt_state, val = step(i, opt_state)
        if i % log_rate == 0 or i == n_iters:
            tr.set_postfix({"Objective": jnp.round(val, 2)})
    return get_params(opt_state)


def optax_fit(
    objective,
    params: dict,
    optax_optim,
    n_iters: int = 100,
    log_rate: int = 10,
) -> tp.Dict:
    opt_state = optax_optim.init(params)
    ps = params

    def step(params, opt_state):
        v, g = jax.value_and_grad(objective)(params)
        updates, opt_state = optax_optim.update(g, opt_state, params)
        params = optax.apply_updates(params, updates)
        return params, opt_state, v

    tr = trange(n_iters)
    for i in tr:
        params, opt_state, val = step(params, opt_state)
        if i % log_rate == 0 or i == n_iters:
            tr.set_postfix({"Objective": jnp.round(val, 2)})
    return params

gpjax/__init__.py METASEP
from jax.config import config

# Enable Floa64 - this is crucial for more stable matrix inversions.
config.update("jax_enable_x64", True)
# Highlight any potentially unintended broadcasting rank promoting ops.
# config.update("jax_numpy_rank_promotion", "warn")

from .abstractions import fit, optax_fit
from .gps import Prior, construct_posterior
from .kernels import (
    RBF,
    Matern12,
    Matern32,
    Matern52,
    Polynomial,
    ProductKernel,
    SumKernel,
)
from .likelihoods import Bernoulli, Gaussian
from .mean_functions import Constant, Zero
from .parameters import copy_dict_structure, initialise, transform
from .types import Dataset

__version__ = "0.4.1"

docs/nbs/datasets/spatial/London_Ward.shx METASEP

docs/nbs/datasets/spatial/London_Ward.shp.xml METASEP
<?xml version="1.0" encoding="UTF-8"?>
<metadata xml:lang="en"><Esri><CreaDate>20140527</CreaDate><CreaTime>17574800</CreaTime><ArcGISFormat>1.0</ArcGISFormat><ArcGISstyle>FGDC CSDGM Metadata</ArcGISstyle><SyncOnce>TRUE</SyncOnce><DataProperties><lineage><Process ToolSource="C:\Program Files (x86)\ArcGIS\Desktop10.0\ArcToolbox\Toolboxes\Data Management Tools.tbx\AddIndex" Date="20140527" Time="175748">AddIndex Q:\Teams\GIS&amp;I\GIS\Processing\OrdnanceSurvey\BoundaryLine_2014_05\ESRI\London\London_Ward.shp GSS_CODE # UNIQUE ASCENDING</Process><Process ToolSource="C:\Program Files (x86)\ArcGIS\Desktop10.0\ArcToolbox\Toolboxes\Data Management Tools.tbx\AddIndex" Date="20140527" Time="175748">AddIndex Q:\Teams\GIS&amp;I\GIS\Processing\OrdnanceSurvey\BoundaryLine_2014_05\ESRI\London\London_Ward.shp LB_GSS_CD # NON_UNIQUE ASCENDING</Process><Process ToolSource="C:\Program Files (x86)\ArcGIS\Desktop10.0\ArcToolbox\Toolboxes\Data Management Tools.tbx\AddIndex" Date="20140527" Time="175749">AddIndex Q:\Teams\GIS&amp;I\GIS\Processing\OrdnanceSurvey\BoundaryLine_2014_05\ESRI\London\London_Ward.shp BOROUGH # NON_UNIQUE ASCENDING</Process><Process ToolSource="C:\Program Files (x86)\ArcGIS\Desktop10.0\ArcToolbox\Toolboxes\Data Management Tools.tbx\AddSpatialIndex" Date="20140527" Time="175753">AddSpatialIndex Q:\Teams\GIS&amp;I\GIS\Processing\OrdnanceSurvey\BoundaryLine_2014_05\ESRI\London\London_Ward.shp 0 0 0</Process></lineage></DataProperties></Esri></metadata>

docs/nbs/datasets/spatial/London_Ward.sbx METASEP

docs/nbs/datasets/spatial/London_Ward.sbn METASEP

docs/nbs/datasets/spatial/London_Ward.prj METASEP
PROJCS["British_National_Grid",GEOGCS["GCS_OSGB_1936",DATUM["D_OSGB_1936",SPHEROID["Airy_1830",6377563.396,299.3249646]],PRIMEM["Greenwich",0],UNIT["Degree",0.0174532925199433]],PROJECTION["Transverse_Mercator"],PARAMETER["False_Easting",400000],PARAMETER["False_Northing",-100000],PARAMETER["Central_Meridian",-2],PARAMETER["Scale_Factor",0.999601272],PARAMETER["Latitude_Of_Origin",49],UNIT["Meter",1]]
docs/nbs/datasets/spatial/London_Ward.dbf METASEP

docs/nbs/datasets/spatial/London_Ward.cpg METASEP
UTF-8
docs/nbs/datasets/spatial/London_Ward.LB_GSS_CD.atx METASEP

docs/nbs/datasets/spatial/London_Ward.GSS_CODE.atx METASEP

docs/nbs/datasets/spatial/London_Ward.BOROUGH.atx METASEP

docs/nbs/datasets/spatial/.DS_Store METASEP

docs/nbs/datasets/aq/aq.shx METASEP

docs/nbs/datasets/aq/aq.shp METASEP

docs/nbs/datasets/aq/aq.dbf METASEP

docs/nbs/datasets/aq/aq.cpg METASEP
ISO-8859-1
setup.cfg METASEP
[bumpversion]
current_version = 0.3.8

[pytest]
testpaths = tests/

[flake8]
max-line-length = 80
max-complexity = 12
ignore = E501
select = C,E,F,W,B,B9
extend-ignore = E203, W503

[tool:mypy]
warn_unused_ignores = True
warn_redundant_casts = True
ignore_missing_imports = True

[mypy-pdoc.test.example_pkg.*]
ignore_errors = True

[tool:pytest]
testpaths = tests/

[tool:pydocstyle]
convention = google

[isort]
profile = black

[metadata]
license_files = LICENSE.txt

requirements.txt METASEP
jax>=0.2.21
optax>=0.1.0
chex>=0.0.8
tensorflow-probability==0.15.0
tqdm
ml-collections==0.1.0
pytest.ini METASEP
[pytest]
testpaths = tests/

codecov.yml METASEP
coverage:
  range: 70..90
  status:
    project:
      default:
        target: auto
        threshold: 5%

fixes:
  - "/home/runner/work/gpjax/gpjax/::" # Correct paths

ignore:
  - "setup.py"
  - "*/tests/.*"
  - "__init__.py"
  - "tests/*.py"

README.md METASEP
# GPJax

[![codecov](https://codecov.io/gh/thomaspinder/gpjax/branch/master/graph/badge.svg?token=DM1DRDASU2)](https://codecov.io/gh/thomaspinder/gpjax)
[![CodeFactor](https://www.codefactor.io/repository/github/thomaspinder/gpjax/badge)](https://www.codefactor.io/repository/github/thomaspinder/gpjax)
[![Documentation Status](https://readthedocs.org/projects/gpjax/badge/?version=latest)](https://gpjax.readthedocs.io/en/latest/?badge=latest)

[**Quickstart**](#simple-example)
| [**Install guide**](#installation)
| [**Documentation**](https://gpjax.readthedocs.io/en/latest/)

GPJax aims to provide a low-level interface to Gaussian process models. Code is written entirely in [Jax](https://github.com/google/jax) to enhance readability, and structured to allow researchers to easily extend the code to suit their own needs. When defining GP prior in GPJax, the user need only specify a mean and kernel function. A GP posterior can then be realised by computing the product of our prior with a likelihood function. The idea behind this is that the code should be as close as possible to the maths that we would write on paper when working with GP models.

## Supported methods and interfaces

### Examples

- [**Conjugate Inference**](https://gpjax.readthedocs.io/en/latest/nbs/regression.html)
- [**Classification**](https://gpjax.readthedocs.io/en/latest/nbs/classification.html)
- [**TensorFlow Probability Integration**](https://gpjax.readthedocs.io/en/latest/nbs/tfp_intergation.html)
- [**Inference on Non-Euclidean Spaces**](https://gpjax.readthedocs.io/en/latest/nbs/kernels.html#Custom-Kernel)

<!-- ### Guides for customisation

- [**Custom Kernel Implementation**](https://gpjax.readthedocs.io/en/latest/nbs/custom_kernel.html) -->

## Simple example

After importing the necessary dependencies, we'll first simulate some data.

```python
import gpjax as gpx
import jax
import jax.numpy as jnp
import jax.random as jr
from jax.experimental import optimizers

key = jr.PRNGKey(123)

x = jr.uniform(key=key, minval=-3.0, maxval=3.0, shape=(50,)).sort().reshape(-1, 1)
y = jnp.sin(x) + jr.normal(key, shape=x.shape)*0.05
training = gpx.Dataset(X=x, y=y)
```

As can be seen, the latent function of interest here is a sinusoidal function. However, it has been perturbed by some zero-mean Gaussian noise with variance of 0.05. We can use a Gaussian process model to try and recover this latent function.

```python
f = gpx.Prior(kernel = gpx.RBF())
```

In the presence of a likelihood function which we'll here assume to be Gaussian, we can optimise the marginal log-likelihood of the Gaussian process prior multiplied by the likelihood to obtain a posterior distribution over the latent function.

```python
likelihood = gpx.Gaussian(num_datapoints = x.shape[0])
posterior = f * likelihood
```

Equipped with the Gaussian process posterior, we can now optimise the model's hyperparameters (note, we need not optimise the latent function here due to the Gaussian conjugacy.). To do this, we can either define our parameters by hand through a dictionary, or realise a set of default parameters through the `initialise` callable. For brevity, we'll do the latter here but see the [regression notebook](https://github.com/thomaspinder/GPJax/blob/master/docs/nbs/regression.ipynb) for a full discussion on parameter initialisation and transformation.

```python
params, constrainer, unconstrainer = gpx.initialise(posterior)
params = gpx.transform(params, unconstrainer)
```

With initial values defined, we can now optimise the hyperparameters' value by carrying out gradient-based optimisation with respect to the GP's marginal log-likelihood. We'll do this now using Jax's built in optimisers, namely the Adam optimiser with a step-size of 0.01. We can also Jit compile our objective function to accelerate training. You'll notice that it is only now that we have incorporated any data into our GP. This is desirable, as this is exactly how model building works in principle too, where we first build our prior model, then observe some data and use this data to build a posterior.

```python
mll = jit(posterior.marginal_log_likelihood(training, constrainer, negative=True))

opt_init, opt_update, get_params = optimizers.adam(step_size=0.01)
opt_state = opt_init(params)
def step(i, opt_state):
    p = get_params(opt_state)
    g = jax.grad(mll)(params)
    return opt_update(i, g, opt_state)

for i in range(100):
    opt_state = step(i, opt_state)
```

Our parameters are now optimised. We can retransfrom these back onto the parameter's original constrained space and, using this learned value, query the GP at a set of test points.

```python
final_params = gpx.transform(get_params(opt_state), constrainer)

xtest = jnp.linspace(-3., 3., 100).reshape(-1, 1)

predictive_mean = posterior.mean(training, final_params)(xtest)
predictive_variance = jnp.diag(posterior.variance(training, final_params)(xtest))
```

## Installation

### Stable version

To install the latest stable version of gpjax run

```bash
pip install gpjax
```

### Development version

To install the lastest, possibly unstable, version, the following steps should be followed. It is by no means compulsory, but we do advise that you do all of the below inside a virtual environment.

```bash
git clone https://github.com/thomaspinder/GPJax.git
cd GPJax
python setup.py develop
```

It is then recommended that you check your installation using the supplied unit tests

```python
python -m pytest tests/
```

Makefile METASEP
.PHONY: conda format style types black test link check notebooks
.DEFAULT_GOAL = help

PYTHON = python
VERSION = 3.8
NAME = py_name
ROOT = ./
PIP = pip
CONDA = conda
SHELL = bash
PKGROOT = gpjax

help:	## Display this help
		@awk 'BEGIN {FS = ":.*##"; printf "\nUsage:\n  make \033[36m\033[0m\n"} /^[a-zA-Z_-]+:.*?##/ { printf "  \033[36m%-15s\033[0m %s\n", $$1, $$2 } /^##@/ { printf "\n\033[1m%s\033[0m\n", substr($$0, 5) } ' $(MAKEFILE_LIST)


##@ Formatting
black:  ## Format code in-place using black.
		black ${PKGROOT}/ tests/ -l 79 .

isort:  ## Format imports in-place using isort.
		isort ${PKGROOT}/ tests/

format: ## Code styling - black, isort
		black ${PKGROOT}/ tests/ -l 100 .
		@printf "\033[1;34mBlack passes!\033[0m\n\n"
		isort ${PKGROOT}/ tests/
		@printf "\033[1;34misort passes!\033[0m\n\n"

##@ Testing
test:  ## Test code using pytest.
		@printf "\033[1;34mRunning tests with pytest...\033[0m\n\n"
		pytest -v gpjax tests
		@printf "\033[1;34mPyTest passes!\033[0m\n\n"

setup.py METASEP
from setuptools import setup, find_packages


def parse_requirements_file(filename):
    with open(filename, encoding="utf-8") as fid:
        requires = [l.strip() for l in fid.readlines() if l]
    return requires


# Optional Packages
EXTRAS = {
    "dev": [
        "black",
        "isort",
        "pylint",
        "flake8",
    ],
    "tests": [
        "pytest",
    ],
    "docs": [
        "furo==2020.12.30b24",
        "nbsphinx==0.8.1",
        "nb-black==1.0.7",
        "matplotlib==3.3.3",
        "sphinx-copybutton==0.3.5",
        "gpviz==0.0.1",
    ],
}

setup(
    name="GPJax",
    version="0.4.1",
    author="Thomas Pinder",
    author_email="t.pinder2@lancaster.ac.uk",
    packages=find_packages(".", exclude=["tests"]),
    license="LICENSE",
    description="Didactic Gaussian processes in Jax.",
    long_description=(
        "GPJax aims to provide a low-level interface to Gaussian process"
        " models. Code is written entirely in Jax to enhance readability, and"
        " structured so as to allow researchers to easily extend the code to"
        " suit their own needs. When defining GP prior in GPJax, the user need"
        " only specify a mean and kernel function. A GP posterior can then be"
        " realised by computing the product of our prior with a likelihood"
        " function. The idea behind this is that the code should be as close as"
        " possible to the maths that we would write on paper when working with"
        " GP models."
    ),
    install_requires=parse_requirements_file("requirements.txt"),
    extras_require=EXTRAS,
    keywords=["gaussian-processes jax machine-learning bayesian"],
)

docs/nbs/datasets/.DS_Store METASEP

docs/requirements.txt METASEP
furo==2020.12.30b24
nbsphinx==0.8.1
nb-black==1.0.7
matplotlib==3.3.3
sphinx-copybutton==0.3.1
docs/make.bat METASEP
@ECHO OFF

pushd %~dp0

REM Command file for Sphinx documentation

if "%SPHINXBUILD%" == "" (
	set SPHINXBUILD=sphinx-build
)
set SOURCEDIR=.
set BUILDDIR=_build

if "%1" == "" goto help

%SPHINXBUILD% >NUL 2>NUL
if errorlevel 9009 (
	echo.
	echo.The 'sphinx-build' command was not found. Make sure you have Sphinx
	echo.installed, then set the SPHINXBUILD environment variable to point
	echo.to the full path of the 'sphinx-build' executable. Alternatively you
	echo.may add the Sphinx directory to PATH.
	echo.
	echo.If you don't have Sphinx installed, grab it from
	echo.http://sphinx-doc.org/
	exit /b 1
)

%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
goto end

:help
%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%

:end
popd

docs/installation.rst METASEP
Installation
======================

Stable version
-----------------

To install the latest stable release using :code:`pip` execute the following:

.. code-block:: bash

    pip install gpjax

GPU support
^^^^^^^^^^^^^^^^^^^

GPU support is enabled through proper configuration of the underlying `Jax <https://github.com/google/jax>`_ installation. CPU enabled forms of both packages are installed as part of the GPJax installation. For GPU Jax support, the following command should be run

.. code-block:: bash

    # Specify your installed CUDA version.
    CUDA_VERSION=11.0
    pip install jaxlib

Then, within a Python shell run

.. code-block:: python

    import jaxlib
    print(jaxlib.__version__)



Development version
--------------------

To install the latest development version of GPJax, run the following set of commands:

.. code-block:: bash

    git clone https://github.com/thomaspinder/GPJax.git
    cd GPJax
    pip install -r requirements.txt
    python setup.py develop

We recommend that this is done from within a virtual environment.
docs/index.rst METASEP
Welcome to GPJax's documentation!
=================================

GPJax is a didactic Gaussian process that supports GPU acceleration and just-in-time compilation. The intention of GPJax is to provide a framework for researchers to rapidly prototype and develop new Gaussian process methods. We seek to provide this by providing an API that seeks to best represent the underlying mathematics of Gaussian processes.

You can view the source code for GPJax `here on Github <https://github.com/thomaspinder/GPJax>`_

Gaussian process 'Hello World' example
-----------------------------------------------------

To gain an intuition for the exposed API provided by GPJax, a simple example of that derives the Gaussian process posterior for regression can be represented by::

    import gpjax
    import jax.numpy as jnp
    import jax.random as jr
    key = jr.PRNGKey(123)

    kernel = gpjax.kernels.RBF()
    f_prior = gpjax.gps.Prior(kernel = kernel)

    likelihood = gpjax.likelihoods.Gaussian()

    posterior = f_prior * likelihood


.. toctree::
    :maxdepth: 1
    :caption: Getting Started

    installation
    nbs/intro

.. toctree::
    :maxdepth: 1
    :caption: Examples

    nbs/regression
    nbs/classification

.. toctree::
    :maxdepth: 1
    :caption: Guides

    nbs/kernels

.. toctree::
    :maxdepth: 1
    :caption: Experimental

    nbs/tfp_integration

docs/Makefile METASEP
# Minimal makefile for Sphinx documentation
#

# You can set these variables from the command line, and also
# from the environment for the first two.
SPHINXOPTS    ?=
SPHINXBUILD   ?= sphinx-build
SOURCEDIR     = .
BUILDDIR      = _build

# Put it first so that "make" without argument is like "make help".
help:
	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)

.PHONY: help Makefile

# Catch-all target: route all unknown targets to Sphinx using the new
# "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
%: Makefile
	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)

docs/conf.py METASEP
# Configuration file for the Sphinx documentation builder.
#
# This file only contains a selection of the most common options. For a full
# list see the documentation:
# https://www.sphinx-doc.org/en/master/usage/configuration.html

# -- Path setup --------------------------------------------------------------

# If extensions (or modules to document with autodoc) are in another directory,
# add these directories to sys.path here. If the directory is relative to the
# documentation root, use os.path.abspath to make it absolute, like shown here.
#
import os
import sys

# sys.path.insert(0, os.path.abspath('.'))
# sys.path.insert(0, os.path.abspath('../'))
# sys.path.insert(0, os.path.abspath('../../'))
# sys.path.insert(0, os.path.abspath('../../gpjax'))
# sys.path.insert(0, os.path.abspath('../../gpjax/gps'))

sys.path.insert(0, os.path.abspath("../.."))

# -- Project information -----------------------------------------------------

project = "GPJax"
copyright = "2021, Thomas Pinder"
author = "Thomas Pinder"

# The full version, including alpha/beta/rc tags
release = "0.3"


# -- General configuration ---------------------------------------------------

# Add any Sphinx extension module names here, as strings. They can be
# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
# ones.
extensions = [
    "nbsphinx",
    "sphinx.ext.autodoc",
    "sphinx.ext.mathjax",
    "sphinx.ext.napoleon",
    "sphinx.ext.autosummary",
    "sphinx.ext.napoleon",
    "sphinx.ext.viewcode",
    "sphinx.ext.autosectionlabel",
    "sphinx_copybutton",
]

# Add any paths that contain templates here, relative to this directory.
templates_path = ["_templates"]

# List of patterns, relative to source directory, that match files and
# directories to ignore when looking for source files.
# This pattern also affects html_static_path and html_extra_path.
exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]


autosummary_generate = True
napolean_use_rtype = False
autodoc_default_options = {
    "member-order": "bysource",
    "special-members": "__init__, __call__",
    "undoc-members": True,
    "exclude-members": "__weakref__,_abc_impl",
}

# -- Options for HTML output -------------------------------------------------

# The theme to use for HTML and HTML Help pages.  See the documentation for
# a list of builtin themes.
#
html_theme = "furo"

html_theme_options = {
    "light_css_variables": {
        "color-brand-primary": "#B5121B",
        "color-brand-content": "#CC3333",
        "color-admonition-background": "orange",
    },
}

# pygments_style = "sphinx"
# pygments_dark_style = "monokai"


# Add any paths that contain custom static files (such as style sheets) here,
# relative to this directory. They are copied after the builtin static files,
# so a file named "default.css" will overwrite the builtin "default.css".
html_static_path = ["_static"]

docs/nbs/tfp_interface.py METASEP
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.13.4
#   kernelspec:
#     display_name: GPJax
#     language: python
#     name: gpjax
# ---

# %% [markdown] colab={"base_uri": "https://localhost:8080/"} id="mfIyurrKVHro" outputId="0232144a-9ba5-4738-d6dd-8c7beb095266" pycharm={"name": "#%% md\n"}
# # TensorFlow Probability Interface

# %%
# %load_ext lab_black

# %% colab={"base_uri": "https://localhost:8080/"} id="RJvA91AKYj52" outputId="88048f85-742e-4163-8eda-4632bcc998f4"
import gpjax
import gpjax.core as gpx
import gpviz as gpv
import jax.numpy as jnp
import jax.random as jr
import tensorflow_probability.substrates.jax as tfp
import matplotlib.pyplot as plt
from jax import grad, jit

tfd = tfp.distributions
key = jr.PRNGKey(123)
plt.style.use(gpv.__stylesheet__)

# %% [markdown] id="qfU3MNxZkR6W"
# ## Simulate some data

# %% colab={"base_uri": "https://localhost:8080/", "height": 282} id="KBB9v752Yzps" outputId="5e16957a-1a27-4d41-8fb8-3423151e6612"
x = jnp.sort(jr.uniform(key, minval=-5.0, maxval=5.0, shape=(100, 1)), axis=0)
f = lambda x: jnp.sin(jnp.pi * x) / (jnp.pi * x)
y = f(x) + jr.normal(key, shape=x.shape) * 0.05
plt.plot(x, f(x), label="Latent fn")
plt.plot(x, y, "o", label="Observations", alpha=0.6)
plt.legend(loc="best")

# %%
training = gpx.Dataset(X=x, y=y)

# %% [markdown] id="DUAIbmAmkUGn"
# ## Define GP

# %% colab={"base_uri": "https://localhost:8080/"} id="RZLbImDSZBWB" outputId="10be6892-cd65-46d7-dee4-ab5ba1604ae8"
prior = gpx.Prior(kernel=gpx.RBF())
posterior = prior * gpx.Gaussian()
print(posterior)

# %% [markdown] id="j9aKKVP1kcIy"
# ## Initialise parameters
#
# All parameters in this model are constrained to be positive. The softplus transformation will be applied to them to map them onto the entire real line.

# %% colab={"base_uri": "https://localhost:8080/"} id="J4f2GzgfjzVT" outputId="1487a58c-02c6-44d6-8bc5-77caabc7abbe"
params = gpx.initialise(posterior)
configs = gpx.get_defaults()
constrainer, unconstrainer = gpx.build_all_transforms(params.keys(), configs)
params = unconstrainer(params)

# %% [markdown] id="oD1L76w2kzuv"
# ## State priors
#
# We'll also place prior distributions on the constrained value of each of our three parameters.

# %% id="3elDDcZ7kAK4"
priors = {
    "lengthscale": tfd.Gamma(1.0, 1.0),
    "variance": tfd.Gamma(2.0, 2.0),
    "obs_noise": tfd.Gamma(2.0, 2.0),
}

# %% [markdown] id="ScWZ0SmhkW2y"
# ## Define target distribution
#
# The marginal log-likelihood distribution can now be computed from the above GP posterior.

# %% id="EDVEDFgLZmiJ"
mll = gpx.marginal_ll(posterior, transform=constrainer, negative=False)

# %% [markdown] id="3Pre2MSQlAFV"
# ### Check the MLL

# %% colab={"base_uri": "https://localhost:8080/"} id="OY0vJRsukEjr" outputId="8feef5c5-d01d-4cff-a872-dc9c7706bd20"
mll(params, training, priors)

# %% colab={"base_uri": "https://localhost:8080/"} id="y2Q7M_uFkL3j" outputId="746f6eeb-b6ac-4a26-cfec-8f1b93ee8db7"
grad(mll)(params, training, priors)

# %% [markdown] id="rzPxSuJSlDFI"
# ## Setup TFP HMC sampler
#
# To allow our MLL function to accept a Jax array as input, not a set of parameters, we'll define a quick helper function to make an arbitary array to our parameter dictionary.

# %% id="jmebSg83kPWp"
from typing import List


def array_to_dict(varray: jnp.DeviceArray, keys: List):
    pdict = {}
    for val, key in zip(varray, keys):
        pdict[key] = val
    return pdict


# %% [markdown] id="FGn6dDyFlW74"
# We'll also fix all our marginal log-likelihood's arguments except the parameters input.

# %% id="3WK3_VrllURp"
from functools import partial

target_dist = partial(mll, training=training, priors=priors)


# %% id="4ldlTlMQljxi"
def build_log_pi(params, target_fn):
    param_keys = list(params.keys())

    def target(params: jnp.DeviceArray):
        coerced_params = array_to_dict(params, param_keys)
        return target_fn(coerced_params)

    return target


# %% id="ix03mQbmltEO"
log_pi = build_log_pi(params, target_dist)

# %% [markdown] id="mBIa7yMwl16f"
# So now we have a functional representation of our GP posterior's marginal log-likelihood that accepts a single Jax array as input. It's worthwhile here checking that evaluation of our new function agrees with our initial marginal log-likelihood function.

# %% colab={"base_uri": "https://localhost:8080/"} id="v7LlBFNbl04z" outputId="7ffe7c78-a1f7-497c-a49e-369ca97b5234"
initial_params = jnp.array(list(params.values())).squeeze()
print(initial_params)

# %% id="VqLGTRVKmO_B"
assert log_pi(initial_params) == mll(params, training, priors)


# %% [markdown] id="3kDRLnUdmsf1"
# With both the marginal log-likelihood and its respective gradient functions in agreeance, we can go ahead and instantiate the TFP HMC sampler. In the interest of time, we'll just sample 500 times from out targer distribution. However, in practice, you'll want to sample for longer.

# %% id="o5b8DfiemZlu"
def run_chain(key, state):
    kernel = tfp.mcmc.NoUTurnSampler(log_pi, 1e-1)
    return tfp.mcmc.sample_chain(
        500,
        current_state=state,
        kernel=kernel,
        trace_fn=lambda _, results: results.target_log_prob,
        seed=key,
    )


# %% colab={"base_uri": "https://localhost:8080/"} id="K70frhaYm7O1" outputId="bc1f4dce-4daa-45b5-8678-ec1292f027f4"
states, log_probs = jit(run_chain)(key, initial_params)

# %% colab={"base_uri": "https://localhost:8080/"} id="pWvImdz5nFE1" outputId="a5c26b9b-f970-4d16-ae70-9f823e80ff99"
states

# %%
constrained_states = jnp.hstack(
    (
        i.reshape(-1, 1)
        for i in constrainer(array_to_dict(states.T, params.keys())).values()
    )
)

# %% colab={"base_uri": "https://localhost:8080/", "height": 350} id="qOU3XXZevhyw" outputId="d768a7da-e574-4e5d-ef9b-811dd8f94812"
burn_in = 100
thin_factor = 2
fig, ax = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

for i, p, pc, t in zip(range(3), states.T, constrained_states.T, params.keys()):
    ax[0][i].plot(p[burn_in:][::thin_factor])
    ax[0][i].set_title(t.replace("_", " "))
    ax[0][i].set_ylabel("Sampled parameter value (unconstrained)")
    ax[1][i].plot(pc[burn_in:][::thin_factor])
    ax[1][i].set_xlabel("Sample index (post-thinning)")
    ax[1][i].set_ylabel("Sampled parameter value (constrained)")

# %%
# %load_ext watermark
# %watermark -n -u -v -iv -w -a "Thomas Pinder"

docs/nbs/tfp_integration.py METASEP
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.13.4
#   kernelspec:
#     display_name: 'Python 3.9.7 64-bit (''gpjax'': conda)'
#     language: python
#     name: python3
# ---

# %%
import gpjax as gpx
import jax
import jax.numpy as jnp
import jax.random as jr
from jax.experimental import optimizers
import matplotlib.pyplot as plt
from pprint import PrettyPrinter
pp=PrettyPrinter(indent=4)

key = jr.PRNGKey(123)

# %%
N = 50
noise = 0.2

x = jr.uniform(key=key, minval=-3.0, maxval=3.0, shape=(N,)).sort().reshape(-1, 1)
f = lambda x: jnp.sin(4 * x) + jnp.cos(2 * x)
signal = f(x)
y = signal + jr.normal(key, shape=signal.shape) * noise
xtest = jnp.linspace(-3.5, 3.5, 500).reshape(-1, 1)
ytest = f(xtest)

# %%
training = gpx.Dataset(X=x, y=y)
posterior = gpx.Prior(kernel=gpx.RBF()) * gpx.Gaussian(num_datapoints=training.n)

# %%
params, constrainers, unconstrainers = gpx.initialise(posterior)

# %%
import typing as tp

def dict_array_coercion(params) -> tp.Tuple[tp.Callable, tp.Callable]:
    flattened_pytree = jax.tree_util.tree_flatten(params)
    def dict_to_array(parameter_dict) -> jnp.DeviceArray:
        return jax.tree_util.tree_flatten(parameter_dict)
    
    def array_to_dict(parameter_array) -> tp.Dict:
        return jax.tree_util.tree_unflatten(flattened_pytree[1], parameter_array)
    
    return dict_to_array, array_to_dict


# %%
dict_to_array, array_to_dict = dict_array_coercion(params)

# %%
parray = dict_to_array(params)
print(parray)

# %%
array_to_dict(parray)

# %%
param_array = jax.tree_util.tree_flatten(params)[0]
params_reconstruct = jax.tree_util.tree_unflatten(jax.tree_util.tree_flatten(params)[1], param_array)
assert params_reconstruct == params

# %%
import tensorflow_probability.substrates.jax as tfp
tfd = tfp.distributions

priors = gpx.parameters.get_param_template(params)
priors['kernel']['lengthscale'] = tfd.Gamma(concentration=jnp.array(1.), rate=jnp.array(1.))
priors['kernel']['variance'] = tfd.Gamma(concentration=jnp.array(1.), rate=jnp.array(1.))
priors['likelihood']['obs_noise'] = tfd.Gamma(concentration=jnp.array(1.), rate=jnp.array(1.))
priors

# %%
mll = posterior.marginal_log_likelihood(training, constrainers,priors=priors, negative=False)
mll(params)

# %%
from functools import partial 


def build_log_pi(params, target):
    def mapper(param_array, params):
        return jax.tree_util.tree_unflatten(jax.tree_util.tree_flatten(params)[1], param_array)

    mapper_fn = partial(mapper, params=params)
    
    def array_mll(parameter_array):
        parameter_dict = mapper_fn([jnp.array(i) for i in parameter_array])
        return target(parameter_dict)
    return array_mll

mll_array_form = build_log_pi(params, mll)

# %%
jax.tree_util.tree_unflatten(jax.tree_util.tree_flatten(params)[1], param_array)

# %%
jax.grad(mll_array_form)(jax.tree_util.tree_flatten(params)[0])

# %%
mll_array_form(jnp.array(jax.tree_util.tree_flatten(params)[0]))


# %%
def run_chain(key, state):
    kernel = tfp.mcmc.NoUTurnSampler(mll_array_form, 1e-1)
    return tfp.mcmc.sample_chain(
        2000,
        current_state=state,
        kernel=kernel,
        trace_fn=lambda _, results: results.target_log_prob,
        seed=key,
    )


# %%
states, log_probs = jax.jit(run_chain)(key, jnp.array(jax.tree_util.tree_flatten(params)[0]))

# %%
states.shape

# %%
plt.plot(states[500:, 1, :], alpha=0.5)

# %%

docs/nbs/spatial_preprocessing.py METASEP
import geopandas as gpd
import pandas as pd
import datetime as dt
import matplotlib.pyplot as plt


def slice_hour(df: pd.DataFrame, timestamp):
    """
    Extract an hour slice of a spatiotemporal dataset.
    """
    df = df[["latitude", "longitude", "date_utc", "no2_ugm3"]]
    df["timestamp"] = pd.to_datetime(df.date_utc)
    return (
        df[df["timestamp"] == timestamp]
        .dropna()
        .drop(["timestamp", "date_utc"], axis=1)
        .reset_index(drop=True)
    )


if __name__ == "__main__":
    # Download air quality data directly from a given url
    aq_data = pd.read_csv(
        "https://data.london.gov.uk/download/breathe-london-aqmesh-pods/267507cc-9740-4ea7-be05-4d6ae16a5e4a/stationary_data.csv"
    )

    # Slice out 9am on 15th February 2019
    hour_of_interest = dt.datetime(year=2019, month=2, day=15, hour=9)
    spatial_data = slice_hour(aq_data, of_interest)

    spatial_data.to_file("datasets/aq/aq.shp")

docs/nbs/helper_fns.py METASEP
from gpjax.parameters import initialise, build_all_transforms
from gpjax.objectives import marginal_ll
from gpjax.sampling import sample
from jax import jit, value_and_grad
import jax.numpy as jnp
from jax.experimental import optimizers
import matplotlib.pyplot as plt


def fit(posterior, nits, data, configs):
    params = initialise(posterior)
    constrainer, unconstrainer = build_all_transforms(params.keys(), configs)

    mll = jit(marginal_ll(posterior, transform=constrainer, negative=True))

    opt_init, opt_update, get_params = optimizers.adam(step_size=0.05)
    opt_state = opt_init(params)

    def step(i, opt_state):
        p = get_params(opt_state)
        v, g = value_and_grad(mll)(p, data)
        return opt_update(i, g, opt_state), v

    for i in range(nits):
        opt_state, mll_estimate = step(i, opt_state)
    print(f"{posterior.prior.kernel.name} GP's marginal log-likelihood: {mll_estimate: .2f}")

    final_params = constrainer(get_params(opt_state))
    return final_params


def plot(key, rv, query_points, training_data, ax, n_samples=100):
    mu = rv.mean()
    sigma = rv.variance()
    one_stddev = 1.96 * jnp.sqrt(sigma)
    posterior_samples = sample(key, rv, n_samples=n_samples)
    plt.grid(color="#888888")  # Color the grid
    ax.spines["polar"].set_visible(False)  # Show or hide the plot spine
    ax.plot(query_points, posterior_samples.T, alpha=0.1)
    ax.fill_between(
        query_points.squeeze(),
        mu - one_stddev,
        mu + one_stddev,
        alpha=0.3,
        label=r"1 Posterior s.d.",
        color="#B5121B"
    )
    ax.fill_between(
        query_points.squeeze(),
        mu - 3*one_stddev,
        mu + 3*one_stddev,
        alpha=0.15,
        label=r"3 Posterior s.d.",
        color="#B5121B"
    )
    ax.plot(query_points, mu, label='Posterior mean')
    ax.scatter(training_data.X, training_data.y, alpha=1, label="Observations")
    ax.legend()

docs/nbs/kernels.py METASEP

# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
slice_kernel = gpx.RBF(active_dims = [0, 1, 3])

# %% [markdown]
# This will implicitly make the kernel have a single lengthscale per input dimensions. Our kernel is now an ARD kernel:

# %%
print(f"ARD: {slice_kernel.ard}")
print(f"Lengthscales: {slice_kernel.params['lengthscale']}")

# %% [markdown]
# We'll now simulate some data and evaluate the kernel on the previously selected input dimensions

# %%
x_matrix = jr.normal(key, shape = (50, 5))
K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params)
print(K.shape)

# %% [markdown]
# ## Combination of Kernels
#
# The product or sum of two positive define matrices yields a positive definite matrix. Consequently, summing or multiplying sets of kernels is a valid operation and can yield rich kernel functions. In GPJax, sums of kernels can be created by passing a list of kernels into the `SumKernel` object in the following way.

# %%
k1 = gpx.RBF()
k1._params = {'lengthscale': jnp.array(1.), 'variance': jnp.array(10.)}
k2 = gpx.Polynomial()
sum_k = gpx.SumKernel(kernel_set = [k1, k2])

fig, ax = plt.subplots(ncols=3, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(sum_k, x, sum_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])

# %% [markdown]
# Similarily, products of kernels can be created through the `ProductKernel` object.

# %%
k3 = gpx.Matern32()
prod_k = gpx.ProductKernel(kernel_set = [k1, k2, k3])
fig, ax = plt.subplots(ncols=4, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(k3, x, k3.params))
im3=ax[3].matshow(gpx.kernels.gram(prod_k, x, prod_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])
fig.colorbar(im3, ax=ax[3])

# %% [markdown]
# ## Custom Kernel
#
# If one wishes to implement their own kernel then the process is straightforward. There are two steps to achieving this 
# 1. List the kernel's parameters
# 2. Define the kernel's pairwise operation
#
# We'll demonstrate this process now for a ciruclar kernel. This example is an adaption of the excellent guide given in the PYMC3 documentation and we would encourage the curious reader to visit their notebook which can be found [here](https://docs.pymc.io/pymc-examples/examples/gaussian_processes/GP-Circular.html).
#
# ### Circular Kernel
#
# When the underlying space is polar, typical Euclidean kernels such as Matern kernels will be insufficient as the boundary. Circular kernels do not exhibit this behaviour and instead _wrap_ around the boundary points to create a smooth function. Such a kernel was given in [Padonou & Roustant (2015)](https://hal.inria.fr/hal-01119942v1) and for two angles $\theta, \theta'$ can be written as
# $$W_c(\theta, \theta') = \left\lvert \left(1 + \tau \frac{d(\theta, \theta')}{c} \right) \left(1 - \frac{d(\theta, \theta')}{c} \right)^{\tau} \right\rvert \quad \tau \geq 4 \tag{1}.$$
#
# Here we have the parameter $\tau$ which is analogous to a lengthscale for Euclidean stationary kernels as it controls the correlation between pairs of observations. Further, the distance metric $d$ here is an angular distance
# $$d(\theta, \theta') = \lvert (\theta-\theta'+c) \operatorname{mod} 2c - c \rvert.$$
#
# To implement this, one must write the following class 

# %%
from chex import dataclass


def angular_distance(x, y, c):
    return jnp.abs((x - y + c) % (c * 2) - c)
    

@dataclass(repr=False)
class Polar(gpx.kernels.Kernel):
    period: float = 2*jnp.pi

    def __post_init__(self):
        self.c = self.period / 2.0  # in [0, \pi]


    def __call__(self, x: gpx.types.Array, y: gpx.types.Array, params: dict) -> gpx.types.Array:
        tau = params["tau"]
        t = angular_distance(x, y, self.c)
        K = (1 + tau * t / self.c) * jnp.clip(1 - t / self.c, 0, jnp.inf) ** tau
        return K.squeeze()

    @property
    def params(self) -> dict:
        return {'tau': jnp.array([4.])}

# %% [markdown]
# We can unpack this now to make better sense of it. In the kernel's `__init__` function we are simply specifying the length of a single period. Here, the underlying domain is a circle, so one period is simply $2\pi$. Next we define the kernel's `__call__` function which is a direct implementation of Equation 1. Finally, we define the Kernel's parameter property which contains just one value $\tau$. Now $\tau$ is defined to be greater than or equal to 4, so we initialise it at 4. 
#
#
# ### Custom Parameter Bijection
#
# The constraint on $\tau$ will make optimisation challenging and it'd be much more helpful if we could transform $\tau$ to be on the full real line when we perform gradient descent on it. Fortunately, the `add_parameter` function in GPJax means we must simply define the parameter's name and corresponding bijection and everything else will be taken care of. Under the hood, a master configuration object is being updated when we call this function to register this parameter and its corresponding transform.

# %%
from gpjax.config import add_parameter
import tensorflow_probability.substrates.jax.bijectors as tfb

add_parameter('tau',  tfb.Softplus(low=jnp.array(4.0)))

# %% [markdown]
# ### Using our Polar Kernel
#
# With a custom kernel now defined, we can go ahead and fit our new kernel. To do this, we randomly simulate a sequence of points on a circle and fit our GP to this using the above defined circular kernel. For further details on this process, it's advised to see the [Regression notebook](https://gpjax.readthedocs.io/en/latest/nbs/regression.html).

# %%
# Simulate data
angles = jnp.linspace(0, 2 * jnp.pi, num=200).reshape(-1, 1)
n = 20
noise = 0.2

X = jnp.sort(jr.uniform(key, minval=0.0, maxval=jnp.pi * 2, shape=(n, 1)), axis=0)
y = 4 + jnp.cos(2 * X) + jr.normal(key, shape=X.shape) * noise
training = gpx.Dataset(X=X, y=y)

# Define Polar Gaussian process
infile
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
slice_kernel = gpx.RBF(active_dims = [0, 1, 3])

# %% [markdown]
# This will implicitly make the kernel have a single lengthscale per input dimensions. Our kernel is now an ARD kernel:

# %%
print(f"ARD: {slice_kernel.ard}")
print(f"Lengthscales: {slice_kernel.params['lengthscale']}")

# %% [markdown]
# We'll now simulate some data and evaluate the kernel on the previously selected input dimensions

# %%
x_matrix = jr.normal(key, shape = (50, 5))
K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params)
print(K.shape)

# %% [markdown]
# ## Combination of Kernels
#
# The product or sum of two positive define matrices yields a positive definite matrix. Consequently, summing or multiplying sets of kernels is a valid operation and can yield rich kernel functions. In GPJax, sums of kernels can be created by passing a list of kernels into the `SumKernel` object in the following way.

# %%
k1 = gpx.RBF()
k1._params = {'lengthscale': jnp.array(1.), 'variance': jnp.array(10.)}
k2 = gpx.Polynomial()
sum_k = gpx.SumKernel(kernel_set = [k1, k2])

fig, ax = plt.subplots(ncols=3, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(sum_k, x, sum_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])

# %% [markdown]
# Similarily, products of kernels can be created through the `ProductKernel` object.

# %%
k3 = gpx.Matern32()
prod_k = gpx.ProductKernel(kernel_set = [k1, k2, k3])
fig, ax = plt.subplots(ncols=4, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(k3, x, k3.params))
im3=ax[3].matshow(gpx.kernels.gram(prod_k, x, prod_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])
fig.colorbar(im3, ax=ax[3])

# %% [markdown]
# ## Custom Kernel
#
# If one wishes to implement their own kernel then the process is straightforward. There are two steps to achieving this 
# 1. List the kernel's parameters
# 2. Define the kernel's pairwise operation
#
# We'll demonstrate this process now for a ciruclar kernel. This example is an adaption of the excellent guide given in the PYMC3 documentation and we would encourage the curious reader to visit their notebook which can be found [here](https://docs.pymc.io/pymc-examples/examples/gaussian_processes/GP-Circular.html).
#
# ### Circular Kernel
#
# When the underlying space is polar, typical Euclidean kernels such as Matern kernels will be insufficient as the boundary. Circular kernels do not exhibit this behaviour and instead _wrap_ around the boundary points to create a smooth function. Such a kernel was given in [Padonou & Roustant (2015)](https://hal.inria.fr/hal-01119942v1) and for two angles $\theta, \theta'$ can be written as
# $$W_c(\theta, \theta') = \left\lvert \left(1 + \tau \frac{d(\theta, \theta')}{c} \right) \left(1 - \frac{d(\theta, \theta')}{c} \right)^{\tau} \right\rvert \quad \tau \geq 4 \tag{1}.$$
#
# Here we have the parameter $\tau$ which is analogous to a lengthscale for Euclidean stationary kernels as it controls the correlation between pairs of observations. Further, the distance metric $d$ here is an angular distance
# $$d(\theta, \theta') = \lvert (\theta-\theta'+c) \operatorname{mod} 2c - c \rvert.$$
#
# To implement this, one must write the following class 

# %%
from chex import dataclass


def angular_distance(x, y, c):
    return jnp.abs((x - y + c) % (c * 2) - c)
    

@dataclass(repr=False)
class Polar(gpx.kernels.Kernel):
    period: float = 2*jnp.pi

    def __post_init__(self):
        self.c = self.period / 2.0  # in [0, \pi]


    def __call__(self, x: gpx.types.Array, y: gpx.types.Array, params: dict) -> gpx.types.Array:
infile
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
slice_kernel = gpx.RBF(active_dims = [0, 1, 3])

# %% [markdown]
# This will implicitly make the kernel have a single lengthscale per input dimensions. Our kernel is now an ARD kernel:

# %%
print(f"ARD: {slice_kernel.ard}")
print(f"Lengthscales: {slice_kernel.params['lengthscale']}")

# %% [markdown]
# We'll now simulate some data and evaluate the kernel on the previously selected input dimensions

# %%
x_matrix = jr.normal(key, shape = (50, 5))
K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params)
print(K.shape)

# %% [markdown]
# ## Combination of Kernels
#
# The product or sum of two positive define matrices yields a positive definite matrix. Consequently, summing or multiplying sets of kernels is a valid operation and can yield rich kernel functions. In GPJax, sums of kernels can be created by passing a list of kernels into the `SumKernel` object in the following way.

# %%
k1 = gpx.RBF()
k1._params = {'lengthscale': jnp.array(1.), 'variance': jnp.array(10.)}
k2 = gpx.Polynomial()
sum_k = gpx.SumKernel(kernel_set = [k1, k2])

fig, ax = plt.subplots(ncols=3, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(sum_k, x, sum_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])

# %% [markdown]
# Similarily, products of kernels can be created through the `ProductKernel` object.

# %%
k3 = gpx.Matern32()
prod_k = gpx.ProductKernel(kernel_set = [k1, k2, k3])
fig, ax = plt.subplots(ncols=4, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(k3, x, k3.params))
im3=ax[3].matshow(gpx.kernels.gram(prod_k, x, prod_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])
fig.colorbar(im3, ax=ax[3])

# %% [markdown]
# ## Custom Kernel
#
# If one wishes to implement their own kernel then the process is straightforward. There are two steps to achieving this 
# 1. List the kernel's parameters
# 2. Define the kernel's pairwise operation
#
# We'll demonstrate this process now for a ciruclar kernel. This example is an adaption of the excellent guide given in the PYMC3 documentation and we would encourage the curious reader to visit their notebook which can be found [here](https://docs.pymc.io/pymc-examples/examples/gaussian_processes/GP-Circular.html).
#
# ### Circular Kernel
#
# When the underlying space is polar, typical Euclidean kernels such as Matern kernels will be insufficient as the boundary. Circular kernels do not exhibit this behaviour and instead _wrap_ around the boundary points to create a smooth function. Such a kernel was given in [Padonou & Roustant (2015)](https://hal.inria.fr/hal-01119942v1) and for two angles $\theta, \theta'$ can be written as
# $$W_c(\theta, \theta') = \left\lvert \left(1 + \tau \frac{d(\theta, \theta')}{c} \right) \left(1 - \frac{d(\theta, \theta')}{c} \right)^{\tau} \right\rvert \quad \tau \geq 4 \tag{1}.$$
#
# Here we have the parameter $\tau$ which is analogous to a lengthscale for Euclidean stationary kernels as it controls the correlation between pairs of observations. Further, the distance metric $d$ here is an angular distance
# $$d(\theta, \theta') = \lvert (\theta-\theta'+c) \operatorname{mod} 2c - c \rvert.$$
#
# To implement this, one must write the following class 

# %%
from chex import dataclass


def angular_distance(x, y, c):
    return jnp.abs((x - y + c) % (c * 2) - c)
    

@dataclass(repr=False)
class Polar(gpx.kernels.Kernel):
    period: float = 2*jnp.pi

    def __post_init__(self):
        self.c = self.period / 2.0  # in [0, \pi]


    def __call__(self, x: gpx.types.Array, y: gpx.types.Array, params: dict) -> gpx.types.Array:
        tau = params["tau"]
infile
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
inproject
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
inproject
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
inproject
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
inproject
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
slice_kernel = gpx.RBF(active_dims = [0, 1, 3])

# %% [markdown]
# This will implicitly make the kernel have a single lengthscale per input dimensions. Our kernel is now an ARD kernel:

# %%
print(f"ARD: {slice_kernel.ard}")
print(f"Lengthscales: {slice_kernel.params['lengthscale']}")

# %% [markdown]
# We'll now simulate some data and evaluate the kernel on the previously selected input dimensions

# %%
x_matrix = jr.normal(key, shape = (50, 5))
K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params)
print(K.shape)

# %% [markdown]
# ## Combination of Kernels
#
# The product or sum of two positive define matrices yields a positive definite matrix. Consequently, summing or multiplying sets of kernels is a valid operation and can yield rich kernel functions. In GPJax, sums of kernels can be created by passing a list of kernels into the `SumKernel` object in the following way.

# %%
k1 = gpx.RBF()
k1._params = {'lengthscale': jnp.array(1.), 'variance': jnp.array(10.)}
k2 = gpx.Polynomial()
sum_k = gpx.SumKernel(kernel_set = [k1, k2])

fig, ax = plt.subplots(ncols=3, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(sum_k, x, sum_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])

# %% [markdown]
# Similarily, products of kernels can be created through the `ProductKernel` object.

# %%
k3 = gpx.Matern32()
prod_k = gpx.ProductKernel(kernel_set = [k1, k2, k3])
fig, ax = plt.subplots(ncols=4, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(k3, x, k3.params))
im3=ax[3].matshow(gpx.kernels.gram(prod_k, x, prod_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])
fig.colorbar(im3, ax=ax[3])

# %% [markdown]
# ## Custom Kernel
#
# If one wishes to implement their own kernel then the process is straightforward. There are two steps to achieving this 
# 1. List the kernel's parameters
# 2. Define the kernel's pairwise operation
#
# We'll demonstrate this process now for a ciruclar kernel. This example is an adaption of the excellent guide given in the PYMC3 documentation and we would encourage the curious reader to visit their notebook which can be found [here](https://docs.pymc.io/pymc-examples/examples/gaussian_processes/GP-Circular.html).
#
# ### Circular Kernel
#
# When the underlying space is polar, typical Euclidean kernels such as Matern kernels will be insufficient as the boundary. Circular kernels do not exhibit this behaviour and instead _wrap_ around the boundary points to create a smooth function. Such a kernel was given in [Padonou & Roustant (2015)](https://hal.inria.fr/hal-01119942v1) and for two angles $\theta, \theta'$ can be written as
# $$W_c(\theta, \theta') = \left\lvert \left(1 + \tau \frac{d(\theta, \theta')}{c} \right) \left(1 - \frac{d(\theta, \theta')}{c} \right)^{\tau} \right\rvert \quad \tau \geq 4 \tag{1}.$$
#
# Here we have the parameter $\tau$ which is analogous to a lengthscale for Euclidean stationary kernels as it controls the correlation between pairs of observations. Further, the distance metric $d$ here is an angular distance
# $$d(\theta, \theta') = \lvert (\theta-\theta'+c) \operatorname{mod} 2c - c \rvert.$$
#
# To implement this, one must write the following class 

# %%
from chex import dataclass


def angular_distance(x, y, c):
    return jnp.abs((x - y + c) % (c * 2) - c)
    

@dataclass(repr=False)
class Polar(gpx.kernels.Kernel):
    period: float = 2*jnp.pi

    def __post_init__(self):
        self.c = self.period / 2.0  # in [0, \pi]


    def __call__(self, x: gpx.types.Array, y: gpx.types.Array, params: dict) -> gpx.types.Array:
        tau = params["tau"]
        t = angular_distance(x, y, self.c)
        K = (1 + tau * t / self.c) * jnp.clip(1 - t / self.c, 0, jnp.inf) ** tau
        return K.squeeze()

    @property
    def params(self) -> dict:
        return {'tau': jnp.array([4.])}

# %% [markdown]
# We can unpack this now to make better sense of it. In the kernel's `__init__` function we are simply specifying the length of a single period. Here, the underlying domain is a circle, so one period is simply $2\pi$. Next we define the kernel's `__call__` function which is a direct implementation of Equation 1. Finally, we define the Kernel's parameter property which contains just one value $\tau$. Now $\tau$ is defined to be greater than or equal to 4, so we initialise it at 4. 
#
#
# ### Custom Parameter Bijection
#
# The constraint on $\tau$ will make optimisation challenging and it'd be much more helpful if we could transform $\tau$ to be on the full real line when we perform gradient descent on it. Fortunately, the `add_parameter` function in GPJax means we must simply define the parameter's name and corresponding bijection and everything else will be taken care of. Under the hood, a master configuration object is being updated when we call this function to register this parameter and its corresponding transform.

# %%
from gpjax.config import add_parameter
import tensorflow_probability.substrates.jax.bijectors as tfb

inproject
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

inproject
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
slice_kernel = gpx.RBF(active_dims = [0, 1, 3])

# %% [markdown]
# This will implicitly make the kernel have a single lengthscale per input dimensions. Our kernel is now an ARD kernel:

# %%
print(f"ARD: {slice_kernel.ard}")
print(f"Lengthscales: {slice_kernel.params['lengthscale']}")

# %% [markdown]
# We'll now simulate some data and evaluate the kernel on the previously selected input dimensions

# %%
x_matrix = jr.normal(key, shape = (50, 5))
K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params)
print(K.shape)

# %% [markdown]
# ## Combination of Kernels
#
# The product or sum of two positive define matrices yields a positive definite matrix. Consequently, summing or multiplying sets of kernels is a valid operation and can yield rich kernel functions. In GPJax, sums of kernels can be created by passing a list of kernels into the `SumKernel` object in the following way.

# %%
k1 = gpx.RBF()
k1._params = {'lengthscale': jnp.array(1.), 'variance': jnp.array(10.)}
k2 = gpx.Polynomial()
sum_k = gpx.SumKernel(kernel_set = [k1, k2])

fig, ax = plt.subplots(ncols=3, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(sum_k, x, sum_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])

# %% [markdown]
# Similarily, products of kernels can be created through the `ProductKernel` object.

# %%
k3 = gpx.Matern32()
prod_k = gpx.ProductKernel(kernel_set = [k1, k2, k3])
fig, ax = plt.subplots(ncols=4, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(k3, x, k3.params))
im3=ax[3].matshow(gpx.kernels.gram(prod_k, x, prod_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])
fig.colorbar(im3, ax=ax[3])

# %% [markdown]
# ## Custom Kernel
#
# If one wishes to implement their own kernel then the process is straightforward. There are two steps to achieving this 
# 1. List the kernel's parameters
# 2. Define the kernel's pairwise operation
#
# We'll demonstrate this process now for a ciruclar kernel. This example is an adaption of the excellent guide given in the PYMC3 documentation and we would encourage the curious reader to visit their notebook which can be found [here](https://docs.pymc.io/pymc-examples/examples/gaussian_processes/GP-Circular.html).
#
# ### Circular Kernel
#
# When the underlying space is polar, typical Euclidean kernels such as Matern kernels will be insufficient as the boundary. Circular kernels do not exhibit this behaviour and instead _wrap_ around the boundary points to create a smooth function. Such a kernel was given in [Padonou & Roustant (2015)](https://hal.inria.fr/hal-01119942v1) and for two angles $\theta, \theta'$ can be written as
# $$W_c(\theta, \theta') = \left\lvert \left(1 + \tau \frac{d(\theta, \theta')}{c} \right) \left(1 - \frac{d(\theta, \theta')}{c} \right)^{\tau} \right\rvert \quad \tau \geq 4 \tag{1}.$$
#
# Here we have the parameter $\tau$ which is analogous to a lengthscale for Euclidean stationary kernels as it controls the correlation between pairs of observations. Further, the distance metric $d$ here is an angular distance
# $$d(\theta, \theta') = \lvert (\theta-\theta'+c) \operatorname{mod} 2c - c \rvert.$$
#
# To implement this, one must write the following class 

# %%
from chex import dataclass


def angular_distance(x, y, c):
    return jnp.abs((x - y + c) % (c * 2) - c)
    

@dataclass(repr=False)
class Polar(gpx.kernels.Kernel):
    period: float = 2*jnp.pi

    def __post_init__(self):
        self.c = self.period / 2.0  # in [0, \pi]


    def __call__(self, x: gpx.types.Array, y: gpx.types.Array, params: dict) -> gpx.types.Array:
        tau = params["tau"]
        t = angular_distance(x, y, self.c)
        K = (1 + tau * t / self.c) * jnp.clip(1 - t / self.c, 0, jnp.inf) ** tau
        return K.squeeze()

    @property
    def params(self) -> dict:
        return {'tau': jnp.array([4.])}

# %% [markdown]
# We can unpack this now to make better sense of it. In the kernel's `__init__` function we are simply specifying the length of a single period. Here, the underlying domain is a circle, so one period is simply $2\pi$. Next we define the kernel's `__call__` function which is a direct implementation of Equation 1. Finally, we define the Kernel's parameter property which contains just one value $\tau$. Now $\tau$ is defined to be greater than or equal to 4, so we initialise it at 4. 
#
#
# ### Custom Parameter Bijection
#
# The constraint on $\tau$ will make optimisation challenging and it'd be much more helpful if we could transform $\tau$ to be on the full real line when we perform gradient descent on it. Fortunately, the `add_parameter` function in GPJax means we must simply define the parameter's name and corresponding bijection and everything else will be taken care of. Under the hood, a master configuration object is being updated when we call this function to register this parameter and its corresponding transform.

# %%
from gpjax.config import add_parameter
import tensorflow_probability.substrates.jax.bijectors as tfb

add_parameter('tau',  tfb.Softplus(low=jnp.array(4.0)))

# %% [markdown]
# ### Using our Polar Kernel
#
# With a custom kernel now defined, we can go ahead and fit our new kernel. To do this, we randomly simulate a sequence of points on a circle and fit our GP to this using the above defined circular kernel. For further details on this process, it's advised to see the [Regression notebook](https://gpjax.readthedocs.io/en/latest/nbs/regression.html).

# %%
# Simulate data
angles = jnp.linspace(0, 2 * jnp.pi, num=200).reshape(-1, 1)
inproject
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
slice_kernel = gpx.RBF(active_dims = [0, 1, 3])

# %% [markdown]
# This will implicitly make the kernel have a single lengthscale per input dimensions. Our kernel is now an ARD kernel:

# %%
print(f"ARD: {slice_kernel.ard}")
print(f"Lengthscales: {slice_kernel.params['lengthscale']}")

# %% [markdown]
# We'll now simulate some data and evaluate the kernel on the previously selected input dimensions

# %%
x_matrix = jr.normal(key, shape = (50, 5))
K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params)
print(K.shape)

# %% [markdown]
# ## Combination of Kernels
#
# The product or sum of two positive define matrices yields a positive definite matrix. Consequently, summing or multiplying sets of kernels is a valid operation and can yield rich kernel functions. In GPJax, sums of kernels can be created by passing a list of kernels into the `SumKernel` object in the following way.

# %%
k1 = gpx.RBF()
k1._params = {'lengthscale': jnp.array(1.), 'variance': jnp.array(10.)}
k2 = gpx.Polynomial()
sum_k = gpx.SumKernel(kernel_set = [k1, k2])

fig, ax = plt.subplots(ncols=3, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(sum_k, x, sum_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])

# %% [markdown]
# Similarily, products of kernels can be created through the `ProductKernel` object.

# %%
k3 = gpx.Matern32()
prod_k = gpx.ProductKernel(kernel_set = [k1, k2, k3])
fig, ax = plt.subplots(ncols=4, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(k3, x, k3.params))
im3=ax[3].matshow(gpx.kernels.gram(prod_k, x, prod_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])
fig.colorbar(im3, ax=ax[3])

# %% [markdown]
# ## Custom Kernel
#
# If one wishes to implement their own kernel then the process is straightforward. There are two steps to achieving this 
# 1. List the kernel's parameters
# 2. Define the kernel's pairwise operation
#
# We'll demonstrate this process now for a ciruclar kernel. This example is an adaption of the excellent guide given in the PYMC3 documentation and we would encourage the curious reader to visit their notebook which can be found [here](https://docs.pymc.io/pymc-examples/examples/gaussian_processes/GP-Circular.html).
#
# ### Circular Kernel
#
# When the underlying space is polar, typical Euclidean kernels such as Matern kernels will be insufficient as the boundary. Circular kernels do not exhibit this behaviour and instead _wrap_ around the boundary points to create a smooth function. Such a kernel was given in [Padonou & Roustant (2015)](https://hal.inria.fr/hal-01119942v1) and for two angles $\theta, \theta'$ can be written as
# $$W_c(\theta, \theta') = \left\lvert \left(1 + \tau \frac{d(\theta, \theta')}{c} \right) \left(1 - \frac{d(\theta, \theta')}{c} \right)^{\tau} \right\rvert \quad \tau \geq 4 \tag{1}.$$
#
# Here we have the parameter $\tau$ which is analogous to a lengthscale for Euclidean stationary kernels as it controls the correlation between pairs of observations. Further, the distance metric $d$ here is an angular distance
# $$d(\theta, \theta') = \lvert (\theta-\theta'+c) \operatorname{mod} 2c - c \rvert.$$
#
# To implement this, one must write the following class 

# %%
from chex import dataclass


def angular_distance(x, y, c):
    return jnp.abs((x - y + c) % (c * 2) - c)
    

@dataclass(repr=False)
class Polar(gpx.kernels.Kernel):
    period: float = 2*jnp.pi

    def __post_init__(self):
        self.c = self.period / 2.0  # in [0, \pi]


    def __call__(self, x: gpx.types.Array, y: gpx.types.Array, params: dict) -> gpx.types.Array:
        tau = params["tau"]
        t = angular_distance(x, y, self.c)
        K = (1 + tau * t / self.c) * jnp.clip(1 - t / self.c, 0, jnp.inf) ** tau
        return K.squeeze()

    @property
    def params(self) -> dict:
        return {'tau': jnp.array([4.])}

# %% [markdown]
# We can unpack this now to make better sense of it. In the kernel's `__init__` function we are simply specifying the length of a single period. Here, the underlying domain is a circle, so one period is simply $2\pi$. Next we define the kernel's `__call__` function which is a direct implementation of Equation 1. Finally, we define the Kernel's parameter property which contains just one value $\tau$. Now $\tau$ is defined to be greater than or equal to 4, so we initialise it at 4. 
#
#
# ### Custom Parameter Bijection
#
# The constraint on $\tau$ will make optimisation challenging and it'd be much more helpful if we could transform $\tau$ to be on the full real line when we perform gradient descent on it. Fortunately, the `add_parameter` function in GPJax means we must simply define the parameter's name and corresponding bijection and everything else will be taken care of. Under the hood, a master configuration object is being updated when we call this function to register this parameter and its corresponding transform.

# %%
from gpjax.config import add_parameter
import tensorflow_probability.substrates.jax.bijectors as tfb

add_parameter('tau',  tfb.Softplus(low=jnp.array(4.0)))

# %% [markdown]
# ### Using our Polar Kernel
#
# With a custom kernel now defined, we can go ahead and fit our new kernel. To do this, we randomly simulate a sequence of points on a circle and fit our GP to this using the above defined circular kernel. For further details on this process, it's advised to see the [Regression notebook](https://gpjax.readthedocs.io/en/latest/nbs/regression.html).

# %%
# Simulate data
angles = jnp.linspace(0, 2 * jnp.pi, num=200).reshape(-1, 1)
n = 20
noise = 0.2

inproject
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
slice_kernel = gpx.RBF(active_dims = [0, 1, 3])

# %% [markdown]
# This will implicitly make the kernel have a single lengthscale per input dimensions. Our kernel is now an ARD kernel:

# %%
print(f"ARD: {slice_kernel.ard}")
print(f"Lengthscales: {slice_kernel.params['lengthscale']}")

# %% [markdown]
# We'll now simulate some data and evaluate the kernel on the previously selected input dimensions

# %%
x_matrix = jr.normal(key, shape = (50, 5))
K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params)
print(K.shape)

# %% [markdown]
# ## Combination of Kernels
#
# The product or sum of two positive define matrices yields a positive definite matrix. Consequently, summing or multiplying sets of kernels is a valid operation and can yield rich kernel functions. In GPJax, sums of kernels can be created by passing a list of kernels into the `SumKernel` object in the following way.

# %%
k1 = gpx.RBF()
k1._params = {'lengthscale': jnp.array(1.), 'variance': jnp.array(10.)}
k2 = gpx.Polynomial()
sum_k = gpx.SumKernel(kernel_set = [k1, k2])

fig, ax = plt.subplots(ncols=3, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(sum_k, x, sum_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])

# %% [markdown]
# Similarily, products of kernels can be created through the `ProductKernel` object.

# %%
k3 = gpx.Matern32()
prod_k = gpx.ProductKernel(kernel_set = [k1, k2, k3])
fig, ax = plt.subplots(ncols=4, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(k3, x, k3.params))
im3=ax[3].matshow(gpx.kernels.gram(prod_k, x, prod_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])
fig.colorbar(im3, ax=ax[3])

# %% [markdown]
# ## Custom Kernel
#
# If one wishes to implement their own kernel then the process is straightforward. There are two steps to achieving this 
# 1. List the kernel's parameters
# 2. Define the kernel's pairwise operation
#
# We'll demonstrate this process now for a ciruclar kernel. This example is an adaption of the excellent guide given in the PYMC3 documentation and we would encourage the curious reader to visit their notebook which can be found [here](https://docs.pymc.io/pymc-examples/examples/gaussian_processes/GP-Circular.html).
#
# ### Circular Kernel
#
# When the underlying space is polar, typical Euclidean kernels such as Matern kernels will be insufficient as the boundary. Circular kernels do not exhibit this behaviour and instead _wrap_ around the boundary points to create a smooth function. Such a kernel was given in [Padonou & Roustant (2015)](https://hal.inria.fr/hal-01119942v1) and for two angles $\theta, \theta'$ can be written as
# $$W_c(\theta, \theta') = \left\lvert \left(1 + \tau \frac{d(\theta, \theta')}{c} \right) \left(1 - \frac{d(\theta, \theta')}{c} \right)^{\tau} \right\rvert \quad \tau \geq 4 \tag{1}.$$
#
# Here we have the parameter $\tau$ which is analogous to a lengthscale for Euclidean stationary kernels as it controls the correlation between pairs of observations. Further, the distance metric $d$ here is an angular distance
# $$d(\theta, \theta') = \lvert (\theta-\theta'+c) \operatorname{mod} 2c - c \rvert.$$
#
# To implement this, one must write the following class 

# %%
from chex import dataclass


def angular_distance(x, y, c):
    return jnp.abs((x - y + c) % (c * 2) - c)
    

@dataclass(repr=False)
class Polar(gpx.kernels.Kernel):
    period: float = 2*jnp.pi

    def __post_init__(self):
        self.c = self.period / 2.0  # in [0, \pi]


    def __call__(self, x: gpx.types.Array, y: gpx.types.Array, params: dict) -> gpx.types.Array:
        tau = params["tau"]
        t = angular_distance(x, y, self.c)
        K = (1 + tau * t / self.c) * jnp.clip(1 - t / self.c, 0, jnp.inf) ** tau
        return K.squeeze()

    @property
    def params(self) -> dict:
        return {'tau': jnp.array([4.])}

# %% [markdown]
# We can unpack this now to make better sense of it. In the kernel's `__init__` function we are simply specifying the length of a single period. Here, the underlying domain is a circle, so one period is simply $2\pi$. Next we define the kernel's `__call__` function which is a direct implementation of Equation 1. Finally, we define the Kernel's parameter property which contains just one value $\tau$. Now $\tau$ is defined to be greater than or equal to 4, so we initialise it at 4. 
#
#
# ### Custom Parameter Bijection
#
# The constraint on $\tau$ will make optimisation challenging and it'd be much more helpful if we could transform $\tau$ to be on the full real line when we perform gradient descent on it. Fortunately, the `add_parameter` function in GPJax means we must simply define the parameter's name and corresponding bijection and everything else will be taken care of. Under the hood, a master configuration object is being updated when we call this function to register this parameter and its corresponding transform.

# %%
from gpjax.config import add_parameter
import tensorflow_probability.substrates.jax.bijectors as tfb

add_parameter('tau',  tfb.Softplus(low=jnp.array(4.0)))

# %% [markdown]
# ### Using our Polar Kernel
#
# With a custom kernel now defined, we can go ahead and fit our new kernel. To do this, we randomly simulate a sequence of points on a circle and fit our GP to this using the above defined circular kernel. For further details on this process, it's advised to see the [Regression notebook](https://gpjax.readthedocs.io/en/latest/nbs/regression.html).

# %%
# Simulate data
angles = jnp.linspace(0, 2 * jnp.pi, num=200).reshape(-1, 1)
n = 20
noise = 0.2

X = jnp.sort(jr.uniform(key, minval=0.0, maxval=jnp.pi * 2, shape=(n, 1)), axis=0)
y = 4 + jnp.cos(2 * X) + jr.normal(key, shape=X.shape) * noise
inproject
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
inproject
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
slice_kernel = gpx.RBF(active_dims = [0, 1, 3])

# %% [markdown]
# This will implicitly make the kernel have a single lengthscale per input dimensions. Our kernel is now an ARD kernel:

# %%
print(f"ARD: {slice_kernel.ard}")
print(f"Lengthscales: {slice_kernel.params['lengthscale']}")

# %% [markdown]
# We'll now simulate some data and evaluate the kernel on the previously selected input dimensions

# %%
x_matrix = jr.normal(key, shape = (50, 5))
K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params)
print(K.shape)

# %% [markdown]
# ## Combination of Kernels
#
# The product or sum of two positive define matrices yields a positive definite matrix. Consequently, summing or multiplying sets of kernels is a valid operation and can yield rich kernel functions. In GPJax, sums of kernels can be created by passing a list of kernels into the `SumKernel` object in the following way.

# %%
k1 = gpx.RBF()
k1._params = {'lengthscale': jnp.array(1.), 'variance': jnp.array(10.)}
k2 = gpx.Polynomial()
sum_k = gpx.SumKernel(kernel_set = [k1, k2])

fig, ax = plt.subplots(ncols=3, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(sum_k, x, sum_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])

# %% [markdown]
# Similarily, products of kernels can be created through the `ProductKernel` object.

# %%
k3 = gpx.Matern32()
prod_k = gpx.ProductKernel(kernel_set = [k1, k2, k3])
fig, ax = plt.subplots(ncols=4, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(k3, x, k3.params))
im3=ax[3].matshow(gpx.kernels.gram(prod_k, x, prod_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])
fig.colorbar(im3, ax=ax[3])

# %% [markdown]
# ## Custom Kernel
#
# If one wishes to implement their own kernel then the process is straightforward. There are two steps to achieving this 
# 1. List the kernel's parameters
# 2. Define the kernel's pairwise operation
#
# We'll demonstrate this process now for a ciruclar kernel. This example is an adaption of the excellent guide given in the PYMC3 documentation and we would encourage the curious reader to visit their notebook which can be found [here](https://docs.pymc.io/pymc-examples/examples/gaussian_processes/GP-Circular.html).
#
# ### Circular Kernel
#
# When the underlying space is polar, typical Euclidean kernels such as Matern kernels will be insufficient as the boundary. Circular kernels do not exhibit this behaviour and instead _wrap_ around the boundary points to create a smooth function. Such a kernel was given in [Padonou & Roustant (2015)](https://hal.inria.fr/hal-01119942v1) and for two angles $\theta, \theta'$ can be written as
# $$W_c(\theta, \theta') = \left\lvert \left(1 + \tau \frac{d(\theta, \theta')}{c} \right) \left(1 - \frac{d(\theta, \theta')}{c} \right)^{\tau} \right\rvert \quad \tau \geq 4 \tag{1}.$$
#
# Here we have the parameter $\tau$ which is analogous to a lengthscale for Euclidean stationary kernels as it controls the correlation between pairs of observations. Further, the distance metric $d$ here is an angular distance
# $$d(\theta, \theta') = \lvert (\theta-\theta'+c) \operatorname{mod} 2c - c \rvert.$$
#
# To implement this, one must write the following class 

# %%
from chex import dataclass


def angular_distance(x, y, c):
    return jnp.abs((x - y + c) % (c * 2) - c)
    

@dataclass(repr=False)
class Polar(gpx.kernels.Kernel):
    period: float = 2*jnp.pi

    def __post_init__(self):
        self.c = self.period / 2.0  # in [0, \pi]


    def __call__(self, x: gpx.types.Array, y: gpx.types.Array, params: dict) -> gpx.types.Array:
        tau = params["tau"]
        t = angular_distance(x, y, self.c)
        K = (1 + tau * t / self.c) * jnp.clip(1 - t / self.c, 0, jnp.inf) ** tau
        return K.squeeze()

    @property
    def params(self) -> dict:
        return {'tau': jnp.array([4.])}

# %% [markdown]
# We can unpack this now to make better sense of it. In the kernel's `__init__` function we are simply specifying the length of a single period. Here, the underlying domain is a circle, so one period is simply $2\pi$. Next we define the kernel's `__call__` function which is a direct implementation of Equation 1. Finally, we define the Kernel's parameter property which contains just one value $\tau$. Now $\tau$ is defined to be greater than or equal to 4, so we initialise it at 4. 
#
#
# ### Custom Parameter Bijection
#
# The constraint on $\tau$ will make optimisation challenging and it'd be much more helpful if we could transform $\tau$ to be on the full real line when we perform gradient descent on it. Fortunately, the `add_parameter` function in GPJax means we must simply define the parameter's name and corresponding bijection and everything else will be taken care of. Under the hood, a master configuration object is being updated when we call this function to register this parameter and its corresponding transform.

# %%
from gpjax.config import add_parameter
import tensorflow_probability.substrates.jax.bijectors as tfb

add_parameter('tau',  tfb.Softplus(low=jnp.array(4.0)))

# %% [markdown]
# ### Using our Polar Kernel
#
# With a custom kernel now defined, we can go ahead and fit our new kernel. To do this, we randomly simulate a sequence of points on a circle and fit our GP to this using the above defined circular kernel. For further details on this process, it's advised to see the [Regression notebook](https://gpjax.readthedocs.io/en/latest/nbs/regression.html).

# %%
# Simulate data
angles = jnp.linspace(0, 2 * jnp.pi, num=200).reshape(-1, 1)
n = 20
noise = 0.2

X = jnp.sort(jr.uniform(key, minval=0.0, maxval=jnp.pi * 2, shape=(n, 1)), axis=0)
y = 4 + jnp.cos(2 * X) + jr.normal(key, shape=X.shape) * noise
training = gpx.Dataset(X=X, y=y)

# Define Polar Gaussian process
PKern = Polar()
inproject
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
slice_kernel = gpx.RBF(active_dims = [0, 1, 3])

# %% [markdown]
# This will implicitly make the kernel have a single lengthscale per input dimensions. Our kernel is now an ARD kernel:

# %%
print(f"ARD: {slice_kernel.ard}")
print(f"Lengthscales: {slice_kernel.params['lengthscale']}")

# %% [markdown]
# We'll now simulate some data and evaluate the kernel on the previously selected input dimensions

# %%
x_matrix = jr.normal(key, shape = (50, 5))
K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params)
print(K.shape)

# %% [markdown]
# ## Combination of Kernels
#
# The product or sum of two positive define matrices yields a positive definite matrix. Consequently, summing or multiplying sets of kernels is a valid operation and can yield rich kernel functions. In GPJax, sums of kernels can be created by passing a list of kernels into the `SumKernel` object in the following way.

# %%
k1 = gpx.RBF()
k1._params = {'lengthscale': jnp.array(1.), 'variance': jnp.array(10.)}
k2 = gpx.Polynomial()
sum_k = gpx.SumKernel(kernel_set = [k1, k2])

fig, ax = plt.subplots(ncols=3, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(sum_k, x, sum_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])

# %% [markdown]
# Similarily, products of kernels can be created through the `ProductKernel` object.

# %%
k3 = gpx.Matern32()
prod_k = gpx.ProductKernel(kernel_set = [k1, k2, k3])
fig, ax = plt.subplots(ncols=4, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(k3, x, k3.params))
im3=ax[3].matshow(gpx.kernels.gram(prod_k, x, prod_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])
fig.colorbar(im3, ax=ax[3])

# %% [markdown]
# ## Custom Kernel
#
# If one wishes to implement their own kernel then the process is straightforward. There are two steps to achieving this 
# 1. List the kernel's parameters
# 2. Define the kernel's pairwise operation
#
# We'll demonstrate this process now for a ciruclar kernel. This example is an adaption of the excellent guide given in the PYMC3 documentation and we would encourage the curious reader to visit their notebook which can be found [here](https://docs.pymc.io/pymc-examples/examples/gaussian_processes/GP-Circular.html).
#
# ### Circular Kernel
#
# When the underlying space is polar, typical Euclidean kernels such as Matern kernels will be insufficient as the boundary. Circular kernels do not exhibit this behaviour and instead _wrap_ around the boundary points to create a smooth function. Such a kernel was given in [Padonou & Roustant (2015)](https://hal.inria.fr/hal-01119942v1) and for two angles $\theta, \theta'$ can be written as
# $$W_c(\theta, \theta') = \left\lvert \left(1 + \tau \frac{d(\theta, \theta')}{c} \right) \left(1 - \frac{d(\theta, \theta')}{c} \right)^{\tau} \right\rvert \quad \tau \geq 4 \tag{1}.$$
#
# Here we have the parameter $\tau$ which is analogous to a lengthscale for Euclidean stationary kernels as it controls the correlation between pairs of observations. Further, the distance metric $d$ here is an angular distance
# $$d(\theta, \theta') = \lvert (\theta-\theta'+c) \operatorname{mod} 2c - c \rvert.$$
#
# To implement this, one must write the following class 

# %%
from chex import dataclass


def angular_distance(x, y, c):
    return jnp.abs((x - y + c) % (c * 2) - c)
    

@dataclass(repr=False)
class Polar(gpx.kernels.Kernel):
    period: float = 2*jnp.pi

    def __post_init__(self):
        self.c = self.period / 2.0  # in [0, \pi]


    def __call__(self, x: gpx.types.Array, y: gpx.types.Array, params: dict) -> gpx.types.Array:
        tau = params["tau"]
        t = angular_distance(x, y, self.c)
        K = (1 + tau * t / self.c) * jnp.clip(1 - t / self.c, 0, jnp.inf) ** tau
        return K.squeeze()

    @property
    def params(self) -> dict:
        return {'tau': jnp.array([4.])}

# %% [markdown]
# We can unpack this now to make better sense of it. In the kernel's `__init__` function we are simply specifying the length of a single period. Here, the underlying domain is a circle, so one period is simply $2\pi$. Next we define the kernel's `__call__` function which is a direct implementation of Equation 1. Finally, we define the Kernel's parameter property which contains just one value $\tau$. Now $\tau$ is defined to be greater than or equal to 4, so we initialise it at 4. 
#
#
# ### Custom Parameter Bijection
#
# The constraint on $\tau$ will make optimisation challenging and it'd be much more helpful if we could transform $\tau$ to be on the full real line when we perform gradient descent on it. Fortunately, the `add_parameter` function in GPJax means we must simply define the parameter's name and corresponding bijection and everything else will be taken care of. Under the hood, a master configuration object is being updated when we call this function to register this parameter and its corresponding transform.

# %%
from gpjax.config import add_parameter
import tensorflow_probability.substrates.jax.bijectors as tfb

add_parameter('tau',  tfb.Softplus(low=jnp.array(4.0)))

# %% [markdown]
# ### Using our Polar Kernel
#
# With a custom kernel now defined, we can go ahead and fit our new kernel. To do this, we randomly simulate a sequence of points on a circle and fit our GP to this using the above defined circular kernel. For further details on this process, it's advised to see the [Regression notebook](https://gpjax.readthedocs.io/en/latest/nbs/regression.html).

# %%
# Simulate data
angles = jnp.linspace(0, 2 * jnp.pi, num=200).reshape(-1, 1)
n = 20
noise = 0.2

X = jnp.sort(jr.uniform(key, minval=0.0, maxval=jnp.pi * 2, shape=(n, 1)), axis=0)
y = 4 + jnp.cos(2 * X) + jr.normal(key, shape=X.shape) * noise
training = gpx.Dataset(X=X, y=y)

# Define Polar Gaussian process
PKern = Polar()
circlular_posterior = gpx.Prior(kernel=PKern) * gpx.Gaussian(num_datapoints=n)

# Initialise parameters and corresponding transformations
inproject
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
slice_kernel = gpx.RBF(active_dims = [0, 1, 3])

# %% [markdown]
# This will implicitly make the kernel have a single lengthscale per input dimensions. Our kernel is now an ARD kernel:

# %%
print(f"ARD: {slice_kernel.ard}")
print(f"Lengthscales: {slice_kernel.params['lengthscale']}")

# %% [markdown]
# We'll now simulate some data and evaluate the kernel on the previously selected input dimensions

# %%
x_matrix = jr.normal(key, shape = (50, 5))
K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params)
print(K.shape)

# %% [markdown]
# ## Combination of Kernels
#
# The product or sum of two positive define matrices yields a positive definite matrix. Consequently, summing or multiplying sets of kernels is a valid operation and can yield rich kernel functions. In GPJax, sums of kernels can be created by passing a list of kernels into the `SumKernel` object in the following way.

# %%
k1 = gpx.RBF()
k1._params = {'lengthscale': jnp.array(1.), 'variance': jnp.array(10.)}
k2 = gpx.Polynomial()
sum_k = gpx.SumKernel(kernel_set = [k1, k2])

fig, ax = plt.subplots(ncols=3, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(sum_k, x, sum_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])

# %% [markdown]
# Similarily, products of kernels can be created through the `ProductKernel` object.

# %%
k3 = gpx.Matern32()
prod_k = gpx.ProductKernel(kernel_set = [k1, k2, k3])
fig, ax = plt.subplots(ncols=4, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(k3, x, k3.params))
im3=ax[3].matshow(gpx.kernels.gram(prod_k, x, prod_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])
fig.colorbar(im3, ax=ax[3])

# %% [markdown]
# ## Custom Kernel
#
# If one wishes to implement their own kernel then the process is straightforward. There are two steps to achieving this 
# 1. List the kernel's parameters
# 2. Define the kernel's pairwise operation
#
# We'll demonstrate this process now for a ciruclar kernel. This example is an adaption of the excellent guide given in the PYMC3 documentation and we would encourage the curious reader to visit their notebook which can be found [here](https://docs.pymc.io/pymc-examples/examples/gaussian_processes/GP-Circular.html).
#
# ### Circular Kernel
#
# When the underlying space is polar, typical Euclidean kernels such as Matern kernels will be insufficient as the boundary. Circular kernels do not exhibit this behaviour and instead _wrap_ around the boundary points to create a smooth function. Such a kernel was given in [Padonou & Roustant (2015)](https://hal.inria.fr/hal-01119942v1) and for two angles $\theta, \theta'$ can be written as
# $$W_c(\theta, \theta') = \left\lvert \left(1 + \tau \frac{d(\theta, \theta')}{c} \right) \left(1 - \frac{d(\theta, \theta')}{c} \right)^{\tau} \right\rvert \quad \tau \geq 4 \tag{1}.$$
#
# Here we have the parameter $\tau$ which is analogous to a lengthscale for Euclidean stationary kernels as it controls the correlation between pairs of observations. Further, the distance metric $d$ here is an angular distance
# $$d(\theta, \theta') = \lvert (\theta-\theta'+c) \operatorname{mod} 2c - c \rvert.$$
#
# To implement this, one must write the following class 

# %%
from chex import dataclass


def angular_distance(x, y, c):
    return jnp.abs((x - y + c) % (c * 2) - c)
    

@dataclass(repr=False)
class Polar(gpx.kernels.Kernel):
    period: float = 2*jnp.pi

    def __post_init__(self):
        self.c = self.period / 2.0  # in [0, \pi]


    def __call__(self, x: gpx.types.Array, y: gpx.types.Array, params: dict) -> gpx.types.Array:
        tau = params["tau"]
        t = angular_distance(x, y, self.c)
        K = (1 + tau * t / self.c) * jnp.clip(1 - t / self.c, 0, jnp.inf) ** tau
        return K.squeeze()

    @property
    def params(self) -> dict:
        return {'tau': jnp.array([4.])}

# %% [markdown]
# We can unpack this now to make better sense of it. In the kernel's `__init__` function we are simply specifying the length of a single period. Here, the underlying domain is a circle, so one period is simply $2\pi$. Next we define the kernel's `__call__` function which is a direct implementation of Equation 1. Finally, we define the Kernel's parameter property which contains just one value $\tau$. Now $\tau$ is defined to be greater than or equal to 4, so we initialise it at 4. 
#
#
# ### Custom Parameter Bijection
#
# The constraint on $\tau$ will make optimisation challenging and it'd be much more helpful if we could transform $\tau$ to be on the full real line when we perform gradient descent on it. Fortunately, the `add_parameter` function in GPJax means we must simply define the parameter's name and corresponding bijection and everything else will be taken care of. Under the hood, a master configuration object is being updated when we call this function to register this parameter and its corresponding transform.

# %%
from gpjax.config import add_parameter
import tensorflow_probability.substrates.jax.bijectors as tfb

add_parameter('tau',  tfb.Softplus(low=jnp.array(4.0)))

# %% [markdown]
# ### Using our Polar Kernel
#
# With a custom kernel now defined, we can go ahead and fit our new kernel. To do this, we randomly simulate a sequence of points on a circle and fit our GP to this using the above defined circular kernel. For further details on this process, it's advised to see the [Regression notebook](https://gpjax.readthedocs.io/en/latest/nbs/regression.html).

# %%
# Simulate data
angles = jnp.linspace(0, 2 * jnp.pi, num=200).reshape(-1, 1)
n = 20
noise = 0.2

X = jnp.sort(jr.uniform(key, minval=0.0, maxval=jnp.pi * 2, shape=(n, 1)), axis=0)
y = 4 + jnp.cos(2 * X) + jr.normal(key, shape=X.shape) * noise
training = gpx.Dataset(X=X, y=y)

# Define Polar Gaussian process
PKern = Polar()
circlular_posterior = gpx.Prior(kernel=PKern) * gpx.Gaussian(num_datapoints=n)

# Initialise parameters and corresponding transformations
params, constrainer, unconstrainer = gpx.initialise(circlular_posterior)

# Optimise GP's marginal log-likelihood using Adam
inproject
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
slice_kernel = gpx.RBF(active_dims = [0, 1, 3])

# %% [markdown]
# This will implicitly make the kernel have a single lengthscale per input dimensions. Our kernel is now an ARD kernel:

# %%
print(f"ARD: {slice_kernel.ard}")
print(f"Lengthscales: {slice_kernel.params['lengthscale']}")

# %% [markdown]
# We'll now simulate some data and evaluate the kernel on the previously selected input dimensions

# %%
x_matrix = jr.normal(key, shape = (50, 5))
K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params)
print(K.shape)

# %% [markdown]
# ## Combination of Kernels
#
# The product or sum of two positive define matrices yields a positive definite matrix. Consequently, summing or multiplying sets of kernels is a valid operation and can yield rich kernel functions. In GPJax, sums of kernels can be created by passing a list of kernels into the `SumKernel` object in the following way.

# %%
k1 = gpx.RBF()
k1._params = {'lengthscale': jnp.array(1.), 'variance': jnp.array(10.)}
k2 = gpx.Polynomial()
sum_k = gpx.SumKernel(kernel_set = [k1, k2])

fig, ax = plt.subplots(ncols=3, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(sum_k, x, sum_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])

# %% [markdown]
# Similarily, products of kernels can be created through the `ProductKernel` object.

# %%
k3 = gpx.Matern32()
prod_k = gpx.ProductKernel(kernel_set = [k1, k2, k3])
fig, ax = plt.subplots(ncols=4, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(k3, x, k3.params))
im3=ax[3].matshow(gpx.kernels.gram(prod_k, x, prod_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])
fig.colorbar(im3, ax=ax[3])

# %% [markdown]
# ## Custom Kernel
#
# If one wishes to implement their own kernel then the process is straightforward. There are two steps to achieving this 
# 1. List the kernel's parameters
# 2. Define the kernel's pairwise operation
#
# We'll demonstrate this process now for a ciruclar kernel. This example is an adaption of the excellent guide given in the PYMC3 documentation and we would encourage the curious reader to visit their notebook which can be found [here](https://docs.pymc.io/pymc-examples/examples/gaussian_processes/GP-Circular.html).
#
# ### Circular Kernel
#
# When the underlying space is polar, typical Euclidean kernels such as Matern kernels will be insufficient as the boundary. Circular kernels do not exhibit this behaviour and instead _wrap_ around the boundary points to create a smooth function. Such a kernel was given in [Padonou & Roustant (2015)](https://hal.inria.fr/hal-01119942v1) and for two angles $\theta, \theta'$ can be written as
# $$W_c(\theta, \theta') = \left\lvert \left(1 + \tau \frac{d(\theta, \theta')}{c} \right) \left(1 - \frac{d(\theta, \theta')}{c} \right)^{\tau} \right\rvert \quad \tau \geq 4 \tag{1}.$$
#
# Here we have the parameter $\tau$ which is analogous to a lengthscale for Euclidean stationary kernels as it controls the correlation between pairs of observations. Further, the distance metric $d$ here is an angular distance
# $$d(\theta, \theta') = \lvert (\theta-\theta'+c) \operatorname{mod} 2c - c \rvert.$$
#
# To implement this, one must write the following class 

# %%
from chex import dataclass


def angular_distance(x, y, c):
    return jnp.abs((x - y + c) % (c * 2) - c)
    

@dataclass(repr=False)
class Polar(gpx.kernels.Kernel):
    period: float = 2*jnp.pi

    def __post_init__(self):
        self.c = self.period / 2.0  # in [0, \pi]


    def __call__(self, x: gpx.types.Array, y: gpx.types.Array, params: dict) -> gpx.types.Array:
        tau = params["tau"]
        t = angular_distance(x, y, self.c)
        K = (1 + tau * t / self.c) * jnp.clip(1 - t / self.c, 0, jnp.inf) ** tau
        return K.squeeze()

    @property
    def params(self) -> dict:
        return {'tau': jnp.array([4.])}

# %% [markdown]
# We can unpack this now to make better sense of it. In the kernel's `__init__` function we are simply specifying the length of a single period. Here, the underlying domain is a circle, so one period is simply $2\pi$. Next we define the kernel's `__call__` function which is a direct implementation of Equation 1. Finally, we define the Kernel's parameter property which contains just one value $\tau$. Now $\tau$ is defined to be greater than or equal to 4, so we initialise it at 4. 
#
#
# ### Custom Parameter Bijection
#
# The constraint on $\tau$ will make optimisation challenging and it'd be much more helpful if we could transform $\tau$ to be on the full real line when we perform gradient descent on it. Fortunately, the `add_parameter` function in GPJax means we must simply define the parameter's name and corresponding bijection and everything else will be taken care of. Under the hood, a master configuration object is being updated when we call this function to register this parameter and its corresponding transform.

# %%
from gpjax.config import add_parameter
import tensorflow_probability.substrates.jax.bijectors as tfb

add_parameter('tau',  tfb.Softplus(low=jnp.array(4.0)))

# %% [markdown]
# ### Using our Polar Kernel
#
# With a custom kernel now defined, we can go ahead and fit our new kernel. To do this, we randomly simulate a sequence of points on a circle and fit our GP to this using the above defined circular kernel. For further details on this process, it's advised to see the [Regression notebook](https://gpjax.readthedocs.io/en/latest/nbs/regression.html).

# %%
# Simulate data
angles = jnp.linspace(0, 2 * jnp.pi, num=200).reshape(-1, 1)
n = 20
noise = 0.2

X = jnp.sort(jr.uniform(key, minval=0.0, maxval=jnp.pi * 2, shape=(n, 1)), axis=0)
y = 4 + jnp.cos(2 * X) + jr.normal(key, shape=X.shape) * noise
training = gpx.Dataset(X=X, y=y)

# Define Polar Gaussian process
PKern = Polar()
circlular_posterior = gpx.Prior(kernel=PKern) * gpx.Gaussian(num_datapoints=n)

# Initialise parameters and corresponding transformations
params, constrainer, unconstrainer = gpx.initialise(circlular_posterior)

# Optimise GP's marginal log-likelihood using Adam
mll = jit(circlular_posterior.marginal_log_likelihood(training, constrainer, negative=True))
inproject
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
slice_kernel = gpx.RBF(active_dims = [0, 1, 3])

# %% [markdown]
# This will implicitly make the kernel have a single lengthscale per input dimensions. Our kernel is now an ARD kernel:

# %%
print(f"ARD: {slice_kernel.ard}")
print(f"Lengthscales: {slice_kernel.params['lengthscale']}")

# %% [markdown]
# We'll now simulate some data and evaluate the kernel on the previously selected input dimensions

# %%
x_matrix = jr.normal(key, shape = (50, 5))
inproject
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
slice_kernel = gpx.RBF(active_dims = [0, 1, 3])

# %% [markdown]
# This will implicitly make the kernel have a single lengthscale per input dimensions. Our kernel is now an ARD kernel:

# %%
print(f"ARD: {slice_kernel.ard}")
print(f"Lengthscales: {slice_kernel.params['lengthscale']}")

# %% [markdown]
# We'll now simulate some data and evaluate the kernel on the previously selected input dimensions

# %%
x_matrix = jr.normal(key, shape = (50, 5))
K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params)
print(K.shape)

# %% [markdown]
# ## Combination of Kernels
#
# The product or sum of two positive define matrices yields a positive definite matrix. Consequently, summing or multiplying sets of kernels is a valid operation and can yield rich kernel functions. In GPJax, sums of kernels can be created by passing a list of kernels into the `SumKernel` object in the following way.

# %%
k1 = gpx.RBF()
k1._params = {'lengthscale': jnp.array(1.), 'variance': jnp.array(10.)}
k2 = gpx.Polynomial()
sum_k = gpx.SumKernel(kernel_set = [k1, k2])

fig, ax = plt.subplots(ncols=3, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(sum_k, x, sum_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])

# %% [markdown]
# Similarily, products of kernels can be created through the `ProductKernel` object.

# %%
k3 = gpx.Matern32()
prod_k = gpx.ProductKernel(kernel_set = [k1, k2, k3])
fig, ax = plt.subplots(ncols=4, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(k3, x, k3.params))
im3=ax[3].matshow(gpx.kernels.gram(prod_k, x, prod_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])
fig.colorbar(im3, ax=ax[3])

# %% [markdown]
# ## Custom Kernel
#
# If one wishes to implement their own kernel then the process is straightforward. There are two steps to achieving this 
# 1. List the kernel's parameters
# 2. Define the kernel's pairwise operation
#
# We'll demonstrate this process now for a ciruclar kernel. This example is an adaption of the excellent guide given in the PYMC3 documentation and we would encourage the curious reader to visit their notebook which can be found [here](https://docs.pymc.io/pymc-examples/examples/gaussian_processes/GP-Circular.html).
#
# ### Circular Kernel
#
# When the underlying space is polar, typical Euclidean kernels such as Matern kernels will be insufficient as the boundary. Circular kernels do not exhibit this behaviour and instead _wrap_ around the boundary points to create a smooth function. Such a kernel was given in [Padonou & Roustant (2015)](https://hal.inria.fr/hal-01119942v1) and for two angles $\theta, \theta'$ can be written as
# $$W_c(\theta, \theta') = \left\lvert \left(1 + \tau \frac{d(\theta, \theta')}{c} \right) \left(1 - \frac{d(\theta, \theta')}{c} \right)^{\tau} \right\rvert \quad \tau \geq 4 \tag{1}.$$
#
# Here we have the parameter $\tau$ which is analogous to a lengthscale for Euclidean stationary kernels as it controls the correlation between pairs of observations. Further, the distance metric $d$ here is an angular distance
# $$d(\theta, \theta') = \lvert (\theta-\theta'+c) \operatorname{mod} 2c - c \rvert.$$
#
# To implement this, one must write the following class 

# %%
from chex import dataclass


def angular_distance(x, y, c):
    return jnp.abs((x - y + c) % (c * 2) - c)
    

@dataclass(repr=False)
class Polar(gpx.kernels.Kernel):
    period: float = 2*jnp.pi

    def __post_init__(self):
        self.c = self.period / 2.0  # in [0, \pi]


    def __call__(self, x: gpx.types.Array, y: gpx.types.Array, params: dict) -> gpx.types.Array:
        tau = params["tau"]
        t = angular_distance(x, y, self.c)
        K = (1 + tau * t / self.c) * jnp.clip(1 - t / self.c, 0, jnp.inf) ** tau
        return K.squeeze()

    @property
    def params(self) -> dict:
        return {'tau': jnp.array([4.])}

# %% [markdown]
# We can unpack this now to make better sense of it. In the kernel's `__init__` function we are simply specifying the length of a single period. Here, the underlying domain is a circle, so one period is simply $2\pi$. Next we define the kernel's `__call__` function which is a direct implementation of Equation 1. Finally, we define the Kernel's parameter property which contains just one value $\tau$. Now $\tau$ is defined to be greater than or equal to 4, so we initialise it at 4. 
#
#
# ### Custom Parameter Bijection
#
# The constraint on $\tau$ will make optimisation challenging and it'd be much more helpful if we could transform $\tau$ to be on the full real line when we perform gradient descent on it. Fortunately, the `add_parameter` function in GPJax means we must simply define the parameter's name and corresponding bijection and everything else will be taken care of. Under the hood, a master configuration object is being updated when we call this function to register this parameter and its corresponding transform.

# %%
from gpjax.config import add_parameter
import tensorflow_probability.substrates.jax.bijectors as tfb

add_parameter('tau',  tfb.Softplus(low=jnp.array(4.0)))

# %% [markdown]
# ### Using our Polar Kernel
#
# With a custom kernel now defined, we can go ahead and fit our new kernel. To do this, we randomly simulate a sequence of points on a circle and fit our GP to this using the above defined circular kernel. For further details on this process, it's advised to see the [Regression notebook](https://gpjax.readthedocs.io/en/latest/nbs/regression.html).

# %%
# Simulate data
angles = jnp.linspace(0, 2 * jnp.pi, num=200).reshape(-1, 1)
n = 20
noise = 0.2

X = jnp.sort(jr.uniform(key, minval=0.0, maxval=jnp.pi * 2, shape=(n, 1)), axis=0)
y = 4 + jnp.cos(2 * X) + jr.normal(key, shape=X.shape) * noise
training = gpx.Dataset(X=X, y=y)

# Define Polar Gaussian process
PKern = Polar()
circlular_posterior = gpx.Prior(kernel=PKern) * gpx.Gaussian(num_datapoints=n)

# Initialise parameters and corresponding transformations
params, constrainer, unconstrainer = gpx.initialise(circlular_posterior)

# Optimise GP's marginal log-likelihood using Adam
mll = jit(circlular_posterior.marginal_log_likelihood(training, constrainer, negative=True))
learned_params = gpx.optax_fit(mll, params, adam(learning_rate=0.05), n_iters=1000)

# Untransform learned parameters
inproject
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
slice_kernel = gpx.RBF(active_dims = [0, 1, 3])

# %% [markdown]
# This will implicitly make the kernel have a single lengthscale per input dimensions. Our kernel is now an ARD kernel:

# %%
print(f"ARD: {slice_kernel.ard}")
print(f"Lengthscales: {slice_kernel.params['lengthscale']}")

# %% [markdown]
# We'll now simulate some data and evaluate the kernel on the previously selected input dimensions

# %%
x_matrix = jr.normal(key, shape = (50, 5))
K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params)
print(K.shape)

# %% [markdown]
# ## Combination of Kernels
#
# The product or sum of two positive define matrices yields a positive definite matrix. Consequently, summing or multiplying sets of kernels is a valid operation and can yield rich kernel functions. In GPJax, sums of kernels can be created by passing a list of kernels into the `SumKernel` object in the following way.

# %%
k1 = gpx.RBF()
k1._params = {'lengthscale': jnp.array(1.), 'variance': jnp.array(10.)}
k2 = gpx.Polynomial()
sum_k = gpx.SumKernel(kernel_set = [k1, k2])

fig, ax = plt.subplots(ncols=3, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(sum_k, x, sum_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])

# %% [markdown]
# Similarily, products of kernels can be created through the `ProductKernel` object.

# %%
k3 = gpx.Matern32()
prod_k = gpx.ProductKernel(kernel_set = [k1, k2, k3])
fig, ax = plt.subplots(ncols=4, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(k3, x, k3.params))
im3=ax[3].matshow(gpx.kernels.gram(prod_k, x, prod_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])
fig.colorbar(im3, ax=ax[3])

# %% [markdown]
# ## Custom Kernel
#
# If one wishes to implement their own kernel then the process is straightforward. There are two steps to achieving this 
# 1. List the kernel's parameters
# 2. Define the kernel's pairwise operation
#
# We'll demonstrate this process now for a ciruclar kernel. This example is an adaption of the excellent guide given in the PYMC3 documentation and we would encourage the curious reader to visit their notebook which can be found [here](https://docs.pymc.io/pymc-examples/examples/gaussian_processes/GP-Circular.html).
#
# ### Circular Kernel
#
# When the underlying space is polar, typical Euclidean kernels such as Matern kernels will be insufficient as the boundary. Circular kernels do not exhibit this behaviour and instead _wrap_ around the boundary points to create a smooth function. Such a kernel was given in [Padonou & Roustant (2015)](https://hal.inria.fr/hal-01119942v1) and for two angles $\theta, \theta'$ can be written as
# $$W_c(\theta, \theta') = \left\lvert \left(1 + \tau \frac{d(\theta, \theta')}{c} \right) \left(1 - \frac{d(\theta, \theta')}{c} \right)^{\tau} \right\rvert \quad \tau \geq 4 \tag{1}.$$
#
# Here we have the parameter $\tau$ which is analogous to a lengthscale for Euclidean stationary kernels as it controls the correlation between pairs of observations. Further, the distance metric $d$ here is an angular distance
# $$d(\theta, \theta') = \lvert (\theta-\theta'+c) \operatorname{mod} 2c - c \rvert.$$
#
# To implement this, one must write the following class 

# %%
from chex import dataclass


def angular_distance(x, y, c):
    return jnp.abs((x - y + c) % (c * 2) - c)
    

@dataclass(repr=False)
class Polar(gpx.kernels.Kernel):
    period: float = 2*jnp.pi

    def __post_init__(self):
        self.c = self.period / 2.0  # in [0, \pi]


    def __call__(self, x: gpx.types.Array, y: gpx.types.Array, params: dict) -> gpx.types.Array:
        tau = params["tau"]
        t = angular_distance(x, y, self.c)
        K = (1 + tau * t / self.c) * jnp.clip(1 - t / self.c, 0, jnp.inf) ** tau
        return K.squeeze()

    @property
    def params(self) -> dict:
        return {'tau': jnp.array([4.])}

# %% [markdown]
# We can unpack this now to make better sense of it. In the kernel's `__init__` function we are simply specifying the length of a single period. Here, the underlying domain is a circle, so one period is simply $2\pi$. Next we define the kernel's `__call__` function which is a direct implementation of Equation 1. Finally, we define the Kernel's parameter property which contains just one value $\tau$. Now $\tau$ is defined to be greater than or equal to 4, so we initialise it at 4. 
#
#
# ### Custom Parameter Bijection
#
# The constraint on $\tau$ will make optimisation challenging and it'd be much more helpful if we could transform $\tau$ to be on the full real line when we perform gradient descent on it. Fortunately, the `add_parameter` function in GPJax means we must simply define the parameter's name and corresponding bijection and everything else will be taken care of. Under the hood, a master configuration object is being updated when we call this function to register this parameter and its corresponding transform.

# %%
from gpjax.config import add_parameter
import tensorflow_probability.substrates.jax.bijectors as tfb

add_parameter('tau',  tfb.Softplus(low=jnp.array(4.0)))

# %% [markdown]
# ### Using our Polar Kernel
#
# With a custom kernel now defined, we can go ahead and fit our new kernel. To do this, we randomly simulate a sequence of points on a circle and fit our GP to this using the above defined circular kernel. For further details on this process, it's advised to see the [Regression notebook](https://gpjax.readthedocs.io/en/latest/nbs/regression.html).

# %%
# Simulate data
angles = jnp.linspace(0, 2 * jnp.pi, num=200).reshape(-1, 1)
n = 20
noise = 0.2

X = jnp.sort(jr.uniform(key, minval=0.0, maxval=jnp.pi * 2, shape=(n, 1)), axis=0)
y = 4 + jnp.cos(2 * X) + jr.normal(key, shape=X.shape) * noise
training = gpx.Dataset(X=X, y=y)

# Define Polar Gaussian process
PKern = Polar()
circlular_posterior = gpx.Prior(kernel=PKern) * gpx.Gaussian(num_datapoints=n)

# Initialise parameters and corresponding transformations
params, constrainer, unconstrainer = gpx.initialise(circlular_posterior)

# Optimise GP's marginal log-likelihood using Adam
mll = jit(circlular_posterior.marginal_log_likelihood(training, constrainer, negative=True))
learned_params = gpx.optax_fit(mll, params, adam(learning_rate=0.05), n_iters=1000)

# Untransform learned parameters
final_params = gpx.transform(learned_params, constrainer)

# %% [markdown]
# ### Prediction
#
# We'll now query the GP's predictive posterior distribution at a linearly spaced set of test points and plot the results.

# %%
inproject
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
slice_kernel = gpx.RBF(active_dims = [0, 1, 3])

# %% [markdown]
# This will implicitly make the kernel have a single lengthscale per input dimensions. Our kernel is now an ARD kernel:

# %%
print(f"ARD: {slice_kernel.ard}")
print(f"Lengthscales: {slice_kernel.params['lengthscale']}")

# %% [markdown]
# We'll now simulate some data and evaluate the kernel on the previously selected input dimensions

# %%
x_matrix = jr.normal(key, shape = (50, 5))
K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params)
print(K.shape)

# %% [markdown]
# ## Combination of Kernels
#
# The product or sum of two positive define matrices yields a positive definite matrix. Consequently, summing or multiplying sets of kernels is a valid operation and can yield rich kernel functions. In GPJax, sums of kernels can be created by passing a list of kernels into the `SumKernel` object in the following way.

# %%
inproject
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
slice_kernel = gpx.RBF(active_dims = [0, 1, 3])

# %% [markdown]
# This will implicitly make the kernel have a single lengthscale per input dimensions. Our kernel is now an ARD kernel:

# %%
print(f"ARD: {slice_kernel.ard}")
print(f"Lengthscales: {slice_kernel.params['lengthscale']}")

# %% [markdown]
# We'll now simulate some data and evaluate the kernel on the previously selected input dimensions

# %%
x_matrix = jr.normal(key, shape = (50, 5))
K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params)
print(K.shape)

# %% [markdown]
# ## Combination of Kernels
#
# The product or sum of two positive define matrices yields a positive definite matrix. Consequently, summing or multiplying sets of kernels is a valid operation and can yield rich kernel functions. In GPJax, sums of kernels can be created by passing a list of kernels into the `SumKernel` object in the following way.

# %%
k1 = gpx.RBF()
k1._params = {'lengthscale': jnp.array(1.), 'variance': jnp.array(10.)}
k2 = gpx.Polynomial()
sum_k = gpx.SumKernel(kernel_set = [k1, k2])

fig, ax = plt.subplots(ncols=3, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(sum_k, x, sum_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])

# %% [markdown]
# Similarily, products of kernels can be created through the `ProductKernel` object.

# %%
k3 = gpx.Matern32()
prod_k = gpx.ProductKernel(kernel_set = [k1, k2, k3])
fig, ax = plt.subplots(ncols=4, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(k3, x, k3.params))
im3=ax[3].matshow(gpx.kernels.gram(prod_k, x, prod_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])
fig.colorbar(im3, ax=ax[3])

# %% [markdown]
# ## Custom Kernel
#
# If one wishes to implement their own kernel then the process is straightforward. There are two steps to achieving this 
# 1. List the kernel's parameters
# 2. Define the kernel's pairwise operation
#
# We'll demonstrate this process now for a ciruclar kernel. This example is an adaption of the excellent guide given in the PYMC3 documentation and we would encourage the curious reader to visit their notebook which can be found [here](https://docs.pymc.io/pymc-examples/examples/gaussian_processes/GP-Circular.html).
#
# ### Circular Kernel
#
# When the underlying space is polar, typical Euclidean kernels such as Matern kernels will be insufficient as the boundary. Circular kernels do not exhibit this behaviour and instead _wrap_ around the boundary points to create a smooth function. Such a kernel was given in [Padonou & Roustant (2015)](https://hal.inria.fr/hal-01119942v1) and for two angles $\theta, \theta'$ can be written as
# $$W_c(\theta, \theta') = \left\lvert \left(1 + \tau \frac{d(\theta, \theta')}{c} \right) \left(1 - \frac{d(\theta, \theta')}{c} \right)^{\tau} \right\rvert \quad \tau \geq 4 \tag{1}.$$
#
# Here we have the parameter $\tau$ which is analogous to a lengthscale for Euclidean stationary kernels as it controls the correlation between pairs of observations. Further, the distance metric $d$ here is an angular distance
# $$d(\theta, \theta') = \lvert (\theta-\theta'+c) \operatorname{mod} 2c - c \rvert.$$
#
# To implement this, one must write the following class 

# %%
from chex import dataclass


def angular_distance(x, y, c):
    return jnp.abs((x - y + c) % (c * 2) - c)
    

@dataclass(repr=False)
class Polar(gpx.kernels.Kernel):
    period: float = 2*jnp.pi

    def __post_init__(self):
        self.c = self.period / 2.0  # in [0, \pi]


    def __call__(self, x: gpx.types.Array, y: gpx.types.Array, params: dict) -> gpx.types.Array:
        tau = params["tau"]
        t = angular_distance(x, y, self.c)
        K = (1 + tau * t / self.c) * jnp.clip(1 - t / self.c, 0, jnp.inf) ** tau
        return K.squeeze()

    @property
    def params(self) -> dict:
        return {'tau': jnp.array([4.])}

# %% [markdown]
# We can unpack this now to make better sense of it. In the kernel's `__init__` function we are simply specifying the length of a single period. Here, the underlying domain is a circle, so one period is simply $2\pi$. Next we define the kernel's `__call__` function which is a direct implementation of Equation 1. Finally, we define the Kernel's parameter property which contains just one value $\tau$. Now $\tau$ is defined to be greater than or equal to 4, so we initialise it at 4. 
#
#
# ### Custom Parameter Bijection
#
# The constraint on $\tau$ will make optimisation challenging and it'd be much more helpful if we could transform $\tau$ to be on the full real line when we perform gradient descent on it. Fortunately, the `add_parameter` function in GPJax means we must simply define the parameter's name and corresponding bijection and everything else will be taken care of. Under the hood, a master configuration object is being updated when we call this function to register this parameter and its corresponding transform.

# %%
from gpjax.config import add_parameter
import tensorflow_probability.substrates.jax.bijectors as tfb

add_parameter('tau',  tfb.Softplus(low=jnp.array(4.0)))

# %% [markdown]
# ### Using our Polar Kernel
#
# With a custom kernel now defined, we can go ahead and fit our new kernel. To do this, we randomly simulate a sequence of points on a circle and fit our GP to this using the above defined circular kernel. For further details on this process, it's advised to see the [Regression notebook](https://gpjax.readthedocs.io/en/latest/nbs/regression.html).

# %%
# Simulate data
angles = jnp.linspace(0, 2 * jnp.pi, num=200).reshape(-1, 1)
n = 20
noise = 0.2

X = jnp.sort(jr.uniform(key, minval=0.0, maxval=jnp.pi * 2, shape=(n, 1)), axis=0)
y = 4 + jnp.cos(2 * X) + jr.normal(key, shape=X.shape) * noise
training = gpx.Dataset(X=X, y=y)

# Define Polar Gaussian process
PKern = Polar()
circlular_posterior = gpx.Prior(kernel=PKern) * gpx.Gaussian(num_datapoints=n)

# Initialise parameters and corresponding transformations
params, constrainer, unconstrainer = gpx.initialise(circlular_posterior)

# Optimise GP's marginal log-likelihood using Adam
mll = jit(circlular_posterior.marginal_log_likelihood(training, constrainer, negative=True))
learned_params = gpx.optax_fit(mll, params, adam(learning_rate=0.05), n_iters=1000)

# Untransform learned parameters
final_params = gpx.transform(learned_params, constrainer)

# %% [markdown]
# ### Prediction
#
# We'll now query the GP's predictive posterior distribution at a linearly spaced set of test points and plot the results.

# %%
mu = circlular_posterior.mean(training, final_params)(angles).squeeze()
inproject
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
slice_kernel = gpx.RBF(active_dims = [0, 1, 3])

# %% [markdown]
# This will implicitly make the kernel have a single lengthscale per input dimensions. Our kernel is now an ARD kernel:

# %%
print(f"ARD: {slice_kernel.ard}")
print(f"Lengthscales: {slice_kernel.params['lengthscale']}")

# %% [markdown]
# We'll now simulate some data and evaluate the kernel on the previously selected input dimensions

# %%
x_matrix = jr.normal(key, shape = (50, 5))
K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params)
print(K.shape)

# %% [markdown]
# ## Combination of Kernels
#
# The product or sum of two positive define matrices yields a positive definite matrix. Consequently, summing or multiplying sets of kernels is a valid operation and can yield rich kernel functions. In GPJax, sums of kernels can be created by passing a list of kernels into the `SumKernel` object in the following way.

# %%
k1 = gpx.RBF()
k1._params = {'lengthscale': jnp.array(1.), 'variance': jnp.array(10.)}
inproject
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
slice_kernel = gpx.RBF(active_dims = [0, 1, 3])

# %% [markdown]
# This will implicitly make the kernel have a single lengthscale per input dimensions. Our kernel is now an ARD kernel:

# %%
print(f"ARD: {slice_kernel.ard}")
print(f"Lengthscales: {slice_kernel.params['lengthscale']}")

# %% [markdown]
# We'll now simulate some data and evaluate the kernel on the previously selected input dimensions

# %%
x_matrix = jr.normal(key, shape = (50, 5))
K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params)
print(K.shape)

# %% [markdown]
# ## Combination of Kernels
#
# The product or sum of two positive define matrices yields a positive definite matrix. Consequently, summing or multiplying sets of kernels is a valid operation and can yield rich kernel functions. In GPJax, sums of kernels can be created by passing a list of kernels into the `SumKernel` object in the following way.

# %%
k1 = gpx.RBF()
k1._params = {'lengthscale': jnp.array(1.), 'variance': jnp.array(10.)}
k2 = gpx.Polynomial()
inproject
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
slice_kernel = gpx.RBF(active_dims = [0, 1, 3])

# %% [markdown]
# This will implicitly make the kernel have a single lengthscale per input dimensions. Our kernel is now an ARD kernel:

# %%
print(f"ARD: {slice_kernel.ard}")
print(f"Lengthscales: {slice_kernel.params['lengthscale']}")

# %% [markdown]
# We'll now simulate some data and evaluate the kernel on the previously selected input dimensions

# %%
x_matrix = jr.normal(key, shape = (50, 5))
K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params)
print(K.shape)

# %% [markdown]
# ## Combination of Kernels
#
# The product or sum of two positive define matrices yields a positive definite matrix. Consequently, summing or multiplying sets of kernels is a valid operation and can yield rich kernel functions. In GPJax, sums of kernels can be created by passing a list of kernels into the `SumKernel` object in the following way.

# %%
k1 = gpx.RBF()
k1._params = {'lengthscale': jnp.array(1.), 'variance': jnp.array(10.)}
k2 = gpx.Polynomial()
sum_k = gpx.SumKernel(kernel_set = [k1, k2])

fig, ax = plt.subplots(ncols=3, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(sum_k, x, sum_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])

# %% [markdown]
# Similarily, products of kernels can be created through the `ProductKernel` object.

# %%
k3 = gpx.Matern32()
prod_k = gpx.ProductKernel(kernel_set = [k1, k2, k3])
fig, ax = plt.subplots(ncols=4, figsize=(20, 5))
inproject
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
slice_kernel = gpx.RBF(active_dims = [0, 1, 3])

# %% [markdown]
# This will implicitly make the kernel have a single lengthscale per input dimensions. Our kernel is now an ARD kernel:

# %%
print(f"ARD: {slice_kernel.ard}")
print(f"Lengthscales: {slice_kernel.params['lengthscale']}")

# %% [markdown]
# We'll now simulate some data and evaluate the kernel on the previously selected input dimensions

# %%
x_matrix = jr.normal(key, shape = (50, 5))
K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params)
print(K.shape)

# %% [markdown]
# ## Combination of Kernels
#
# The product or sum of two positive define matrices yields a positive definite matrix. Consequently, summing or multiplying sets of kernels is a valid operation and can yield rich kernel functions. In GPJax, sums of kernels can be created by passing a list of kernels into the `SumKernel` object in the following way.

# %%
k1 = gpx.RBF()
k1._params = {'lengthscale': jnp.array(1.), 'variance': jnp.array(10.)}
k2 = gpx.Polynomial()
sum_k = gpx.SumKernel(kernel_set = [k1, k2])

fig, ax = plt.subplots(ncols=3, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(sum_k, x, sum_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])

# %% [markdown]
# Similarily, products of kernels can be created through the `ProductKernel` object.

# %%
k3 = gpx.Matern32()
prod_k = gpx.ProductKernel(kernel_set = [k1, k2, k3])
fig, ax = plt.subplots(ncols=4, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
inproject
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
slice_kernel = gpx.RBF(active_dims = [0, 1, 3])

# %% [markdown]
# This will implicitly make the kernel have a single lengthscale per input dimensions. Our kernel is now an ARD kernel:

# %%
print(f"ARD: {slice_kernel.ard}")
print(f"Lengthscales: {slice_kernel.params['lengthscale']}")

# %% [markdown]
# We'll now simulate some data and evaluate the kernel on the previously selected input dimensions

# %%
x_matrix = jr.normal(key, shape = (50, 5))
K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params)
print(K.shape)

# %% [markdown]
# ## Combination of Kernels
#
# The product or sum of two positive define matrices yields a positive definite matrix. Consequently, summing or multiplying sets of kernels is a valid operation and can yield rich kernel functions. In GPJax, sums of kernels can be created by passing a list of kernels into the `SumKernel` object in the following way.

# %%
k1 = gpx.RBF()
k1._params = {'lengthscale': jnp.array(1.), 'variance': jnp.array(10.)}
k2 = gpx.Polynomial()
sum_k = gpx.SumKernel(kernel_set = [k1, k2])

fig, ax = plt.subplots(ncols=3, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
inproject
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
slice_kernel = gpx.RBF(active_dims = [0, 1, 3])

# %% [markdown]
# This will implicitly make the kernel have a single lengthscale per input dimensions. Our kernel is now an ARD kernel:

# %%
print(f"ARD: {slice_kernel.ard}")
print(f"Lengthscales: {slice_kernel.params['lengthscale']}")

# %% [markdown]
# We'll now simulate some data and evaluate the kernel on the previously selected input dimensions

# %%
x_matrix = jr.normal(key, shape = (50, 5))
K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params)
print(K.shape)

# %% [markdown]
# ## Combination of Kernels
#
# The product or sum of two positive define matrices yields a positive definite matrix. Consequently, summing or multiplying sets of kernels is a valid operation and can yield rich kernel functions. In GPJax, sums of kernels can be created by passing a list of kernels into the `SumKernel` object in the following way.

# %%
k1 = gpx.RBF()
k1._params = {'lengthscale': jnp.array(1.), 'variance': jnp.array(10.)}
k2 = gpx.Polynomial()
sum_k = gpx.SumKernel(kernel_set = [k1, k2])

fig, ax = plt.subplots(ncols=3, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(sum_k, x, sum_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])

# %% [markdown]
# Similarily, products of kernels can be created through the `ProductKernel` object.

# %%
inproject
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
slice_kernel = gpx.RBF(active_dims = [0, 1, 3])

# %% [markdown]
# This will implicitly make the kernel have a single lengthscale per input dimensions. Our kernel is now an ARD kernel:

# %%
print(f"ARD: {slice_kernel.ard}")
print(f"Lengthscales: {slice_kernel.params['lengthscale']}")

# %% [markdown]
# We'll now simulate some data and evaluate the kernel on the previously selected input dimensions

# %%
x_matrix = jr.normal(key, shape = (50, 5))
K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params)
print(K.shape)

# %% [markdown]
# ## Combination of Kernels
#
# The product or sum of two positive define matrices yields a positive definite matrix. Consequently, summing or multiplying sets of kernels is a valid operation and can yield rich kernel functions. In GPJax, sums of kernels can be created by passing a list of kernels into the `SumKernel` object in the following way.

# %%
k1 = gpx.RBF()
k1._params = {'lengthscale': jnp.array(1.), 'variance': jnp.array(10.)}
k2 = gpx.Polynomial()
sum_k = gpx.SumKernel(kernel_set = [k1, k2])

fig, ax = plt.subplots(ncols=3, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(sum_k, x, sum_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])

# %% [markdown]
# Similarily, products of kernels can be created through the `ProductKernel` object.

# %%
k3 = gpx.Matern32()
inproject
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
slice_kernel = gpx.RBF(active_dims = [0, 1, 3])

# %% [markdown]
# This will implicitly make the kernel have a single lengthscale per input dimensions. Our kernel is now an ARD kernel:

# %%
print(f"ARD: {slice_kernel.ard}")
print(f"Lengthscales: {slice_kernel.params['lengthscale']}")

# %% [markdown]
# We'll now simulate some data and evaluate the kernel on the previously selected input dimensions

# %%
x_matrix = jr.normal(key, shape = (50, 5))
K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params)
print(K.shape)

# %% [markdown]
# ## Combination of Kernels
#
# The product or sum of two positive define matrices yields a positive definite matrix. Consequently, summing or multiplying sets of kernels is a valid operation and can yield rich kernel functions. In GPJax, sums of kernels can be created by passing a list of kernels into the `SumKernel` object in the following way.

# %%
k1 = gpx.RBF()
k1._params = {'lengthscale': jnp.array(1.), 'variance': jnp.array(10.)}
k2 = gpx.Polynomial()
sum_k = gpx.SumKernel(kernel_set = [k1, k2])

fig, ax = plt.subplots(ncols=3, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(sum_k, x, sum_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])

# %% [markdown]
# Similarily, products of kernels can be created through the `ProductKernel` object.

# %%
k3 = gpx.Matern32()
prod_k = gpx.ProductKernel(kernel_set = [k1, k2, k3])
fig, ax = plt.subplots(ncols=4, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
inproject
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
slice_kernel = gpx.RBF(active_dims = [0, 1, 3])

# %% [markdown]
# This will implicitly make the kernel have a single lengthscale per input dimensions. Our kernel is now an ARD kernel:

# %%
print(f"ARD: {slice_kernel.ard}")
print(f"Lengthscales: {slice_kernel.params['lengthscale']}")

# %% [markdown]
# We'll now simulate some data and evaluate the kernel on the previously selected input dimensions

# %%
x_matrix = jr.normal(key, shape = (50, 5))
K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params)
print(K.shape)

# %% [markdown]
# ## Combination of Kernels
#
# The product or sum of two positive define matrices yields a positive definite matrix. Consequently, summing or multiplying sets of kernels is a valid operation and can yield rich kernel functions. In GPJax, sums of kernels can be created by passing a list of kernels into the `SumKernel` object in the following way.

# %%
k1 = gpx.RBF()
k1._params = {'lengthscale': jnp.array(1.), 'variance': jnp.array(10.)}
k2 = gpx.Polynomial()
sum_k = gpx.SumKernel(kernel_set = [k1, k2])

fig, ax = plt.subplots(ncols=3, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(sum_k, x, sum_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])

# %% [markdown]
# Similarily, products of kernels can be created through the `ProductKernel` object.

# %%
k3 = gpx.Matern32()
prod_k = gpx.ProductKernel(kernel_set = [k1, k2, k3])
fig, ax = plt.subplots(ncols=4, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(k3, x, k3.params))
inproject
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
slice_kernel = gpx.RBF(active_dims = [0, 1, 3])

# %% [markdown]
# This will implicitly make the kernel have a single lengthscale per input dimensions. Our kernel is now an ARD kernel:

# %%
print(f"ARD: {slice_kernel.ard}")
print(f"Lengthscales: {slice_kernel.params['lengthscale']}")

# %% [markdown]
# We'll now simulate some data and evaluate the kernel on the previously selected input dimensions

# %%
x_matrix = jr.normal(key, shape = (50, 5))
K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params)
print(K.shape)

# %% [markdown]
# ## Combination of Kernels
#
# The product or sum of two positive define matrices yields a positive definite matrix. Consequently, summing or multiplying sets of kernels is a valid operation and can yield rich kernel functions. In GPJax, sums of kernels can be created by passing a list of kernels into the `SumKernel` object in the following way.

# %%
k1 = gpx.RBF()
k1._params = {'lengthscale': jnp.array(1.), 'variance': jnp.array(10.)}
k2 = gpx.Polynomial()
sum_k = gpx.SumKernel(kernel_set = [k1, k2])

fig, ax = plt.subplots(ncols=3, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(sum_k, x, sum_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])

# %% [markdown]
# Similarily, products of kernels can be created through the `ProductKernel` object.

# %%
k3 = gpx.Matern32()
prod_k = gpx.ProductKernel(kernel_set = [k1, k2, k3])
fig, ax = plt.subplots(ncols=4, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(k3, x, k3.params))
im3=ax[3].matshow(gpx.kernels.gram(prod_k, x, prod_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])
fig.colorbar(im3, ax=ax[3])

# %% [markdown]
# ## Custom Kernel
#
# If one wishes to implement their own kernel then the process is straightforward. There are two steps to achieving this 
# 1. List the kernel's parameters
# 2. Define the kernel's pairwise operation
#
# We'll demonstrate this process now for a ciruclar kernel. This example is an adaption of the excellent guide given in the PYMC3 documentation and we would encourage the curious reader to visit their notebook which can be found [here](https://docs.pymc.io/pymc-examples/examples/gaussian_processes/GP-Circular.html).
#
# ### Circular Kernel
#
# When the underlying space is polar, typical Euclidean kernels such as Matern kernels will be insufficient as the boundary. Circular kernels do not exhibit this behaviour and instead _wrap_ around the boundary points to create a smooth function. Such a kernel was given in [Padonou & Roustant (2015)](https://hal.inria.fr/hal-01119942v1) and for two angles $\theta, \theta'$ can be written as
# $$W_c(\theta, \theta') = \left\lvert \left(1 + \tau \frac{d(\theta, \theta')}{c} \right) \left(1 - \frac{d(\theta, \theta')}{c} \right)^{\tau} \right\rvert \quad \tau \geq 4 \tag{1}.$$
#
# Here we have the parameter $\tau$ which is analogous to a lengthscale for Euclidean stationary kernels as it controls the correlation between pairs of observations. Further, the distance metric $d$ here is an angular distance
# $$d(\theta, \theta') = \lvert (\theta-\theta'+c) \operatorname{mod} 2c - c \rvert.$$
#
# To implement this, one must write the following class 

# %%
from chex import dataclass


def angular_distance(x, y, c):
    return jnp.abs((x - y + c) % (c * 2) - c)
    

@dataclass(repr=False)
class Polar(gpx.kernels.Kernel):
    period: float = 2*jnp.pi

    def __post_init__(self):
        self.c = self.period / 2.0  # in [0, \pi]


    def __call__(self, x: gpx.types.Array, y: gpx.types.Array, params: dict) -> gpx.types.Array:
        tau = params["tau"]
        t = angular_distance(x, y, self.c)
        K = (1 + tau * t / self.c) * jnp.clip(1 - t / self.c, 0, jnp.inf) ** tau
        return K.squeeze()

    @property
    def params(self) -> dict:
        return {'tau': jnp.array([4.])}

# %% [markdown]
# We can unpack this now to make better sense of it. In the kernel's `__init__` function we are simply specifying the length of a single period. Here, the underlying domain is a circle, so one period is simply $2\pi$. Next we define the kernel's `__call__` function which is a direct implementation of Equation 1. Finally, we define the Kernel's parameter property which contains just one value $\tau$. Now $\tau$ is defined to be greater than or equal to 4, so we initialise it at 4. 
#
#
# ### Custom Parameter Bijection
#
# The constraint on $\tau$ will make optimisation challenging and it'd be much more helpful if we could transform $\tau$ to be on the full real line when we perform gradient descent on it. Fortunately, the `add_parameter` function in GPJax means we must simply define the parameter's name and corresponding bijection and everything else will be taken care of. Under the hood, a master configuration object is being updated when we call this function to register this parameter and its corresponding transform.

# %%
from gpjax.config import add_parameter
import tensorflow_probability.substrates.jax.bijectors as tfb

add_parameter('tau',  tfb.Softplus(low=jnp.array(4.0)))

# %% [markdown]
# ### Using our Polar Kernel
#
# With a custom kernel now defined, we can go ahead and fit our new kernel. To do this, we randomly simulate a sequence of points on a circle and fit our GP to this using the above defined circular kernel. For further details on this process, it's advised to see the [Regression notebook](https://gpjax.readthedocs.io/en/latest/nbs/regression.html).

# %%
# Simulate data
angles = jnp.linspace(0, 2 * jnp.pi, num=200).reshape(-1, 1)
n = 20
noise = 0.2

X = jnp.sort(jr.uniform(key, minval=0.0, maxval=jnp.pi * 2, shape=(n, 1)), axis=0)
y = 4 + jnp.cos(2 * X) + jr.normal(key, shape=X.shape) * noise
training = gpx.Dataset(X=X, y=y)

# Define Polar Gaussian process
PKern = Polar()
circlular_posterior = gpx.Prior(kernel=PKern) * gpx.Gaussian(num_datapoints=n)

# Initialise parameters and corresponding transformations
params, constrainer, unconstrainer = gpx.initialise(circlular_posterior)

# Optimise GP's marginal log-likelihood using Adam
mll = jit(circlular_posterior.marginal_log_likelihood(training, constrainer, negative=True))
learned_params = gpx.optax_fit(mll, params, adam(learning_rate=0.05), n_iters=1000)

# Untransform learned parameters
final_params = gpx.transform(learned_params, constrainer)

# %% [markdown]
# ### Prediction
#
# We'll now query the GP's predictive posterior distribution at a linearly spaced set of test points and plot the results.

# %%
mu = circlular_posterior.mean(training, final_params)(angles).squeeze()
one_sigma = jnp.sqrt(jnp.diag(circlular_posterior.variance(training, final_params)(angles)))

# %%
fig = plt.figure(figsize=(10, 8))
gridspec = fig.add_gridspec(1, 1)
ax = plt.subplot(gridspec[0], polar=True)

ax.grid(color="#888888")  # Color the grid
ax.spines["polar"].set_visible(False)  # Show or hide the plot spine
ax.fill_between(
    angles.squeeze(),
    mu - one_sigma,
    mu + one_sigma,
    alpha=0.3,
    label=r"1 Posterior s.d.",
    color="#B5121B"
)
ax.fill_between(
    angles.squeeze(),
    mu - 3*one_sigma,
    mu + 3*one_sigma,
    alpha=0.15,
    label=r"3 Posterior s.d.",
    color="#B5121B"
)
inproject
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
slice_kernel = gpx.RBF(active_dims = [0, 1, 3])

# %% [markdown]
# This will implicitly make the kernel have a single lengthscale per input dimensions. Our kernel is now an ARD kernel:

# %%
print(f"ARD: {slice_kernel.ard}")
print(f"Lengthscales: {slice_kernel.params['lengthscale']}")

# %% [markdown]
# We'll now simulate some data and evaluate the kernel on the previously selected input dimensions

# %%
x_matrix = jr.normal(key, shape = (50, 5))
K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params)
print(K.shape)

# %% [markdown]
# ## Combination of Kernels
#
# The product or sum of two positive define matrices yields a positive definite matrix. Consequently, summing or multiplying sets of kernels is a valid operation and can yield rich kernel functions. In GPJax, sums of kernels can be created by passing a list of kernels into the `SumKernel` object in the following way.

# %%
k1 = gpx.RBF()
k1._params = {'lengthscale': jnp.array(1.), 'variance': jnp.array(10.)}
k2 = gpx.Polynomial()
sum_k = gpx.SumKernel(kernel_set = [k1, k2])

fig, ax = plt.subplots(ncols=3, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(sum_k, x, sum_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])

# %% [markdown]
# Similarily, products of kernels can be created through the `ProductKernel` object.

# %%
k3 = gpx.Matern32()
prod_k = gpx.ProductKernel(kernel_set = [k1, k2, k3])
fig, ax = plt.subplots(ncols=4, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(k3, x, k3.params))
im3=ax[3].matshow(gpx.kernels.gram(prod_k, x, prod_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])
fig.colorbar(im3, ax=ax[3])

# %% [markdown]
# ## Custom Kernel
#
# If one wishes to implement their own kernel then the process is straightforward. There are two steps to achieving this 
# 1. List the kernel's parameters
# 2. Define the kernel's pairwise operation
#
# We'll demonstrate this process now for a ciruclar kernel. This example is an adaption of the excellent guide given in the PYMC3 documentation and we would encourage the curious reader to visit their notebook which can be found [here](https://docs.pymc.io/pymc-examples/examples/gaussian_processes/GP-Circular.html).
#
# ### Circular Kernel
#
# When the underlying space is polar, typical Euclidean kernels such as Matern kernels will be insufficient as the boundary. Circular kernels do not exhibit this behaviour and instead _wrap_ around the boundary points to create a smooth function. Such a kernel was given in [Padonou & Roustant (2015)](https://hal.inria.fr/hal-01119942v1) and for two angles $\theta, \theta'$ can be written as
# $$W_c(\theta, \theta') = \left\lvert \left(1 + \tau \frac{d(\theta, \theta')}{c} \right) \left(1 - \frac{d(\theta, \theta')}{c} \right)^{\tau} \right\rvert \quad \tau \geq 4 \tag{1}.$$
#
# Here we have the parameter $\tau$ which is analogous to a lengthscale for Euclidean stationary kernels as it controls the correlation between pairs of observations. Further, the distance metric $d$ here is an angular distance
# $$d(\theta, \theta') = \lvert (\theta-\theta'+c) \operatorname{mod} 2c - c \rvert.$$
#
# To implement this, one must write the following class 

# %%
from chex import dataclass


def angular_distance(x, y, c):
    return jnp.abs((x - y + c) % (c * 2) - c)
    

@dataclass(repr=False)
class Polar(gpx.kernels.Kernel):
    period: float = 2*jnp.pi

    def __post_init__(self):
        self.c = self.period / 2.0  # in [0, \pi]

non_informative
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
non_informative
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
slice_kernel = gpx.RBF(active_dims = [0, 1, 3])

# %% [markdown]
# This will implicitly make the kernel have a single lengthscale per input dimensions. Our kernel is now an ARD kernel:

# %%
print(f"ARD: {slice_kernel.ard}")
print(f"Lengthscales: {slice_kernel.params['lengthscale']}")

# %% [markdown]
# We'll now simulate some data and evaluate the kernel on the previously selected input dimensions

# %%
x_matrix = jr.normal(key, shape = (50, 5))
K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params)
print(K.shape)

# %% [markdown]
# ## Combination of Kernels
#
# The product or sum of two positive define matrices yields a positive definite matrix. Consequently, summing or multiplying sets of kernels is a valid operation and can yield rich kernel functions. In GPJax, sums of kernels can be created by passing a list of kernels into the `SumKernel` object in the following way.

# %%
k1 = gpx.RBF()
k1._params = {'lengthscale': jnp.array(1.), 'variance': jnp.array(10.)}
k2 = gpx.Polynomial()
sum_k = gpx.SumKernel(kernel_set = [k1, k2])

fig, ax = plt.subplots(ncols=3, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(sum_k, x, sum_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])

# %% [markdown]
# Similarily, products of kernels can be created through the `ProductKernel` object.

# %%
k3 = gpx.Matern32()
prod_k = gpx.ProductKernel(kernel_set = [k1, k2, k3])
fig, ax = plt.subplots(ncols=4, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(k3, x, k3.params))
im3=ax[3].matshow(gpx.kernels.gram(prod_k, x, prod_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])
fig.colorbar(im3, ax=ax[3])

# %% [markdown]
# ## Custom Kernel
#
# If one wishes to implement their own kernel then the process is straightforward. There are two steps to achieving this 
# 1. List the kernel's parameters
# 2. Define the kernel's pairwise operation
#
# We'll demonstrate this process now for a ciruclar kernel. This example is an adaption of the excellent guide given in the PYMC3 documentation and we would encourage the curious reader to visit their notebook which can be found [here](https://docs.pymc.io/pymc-examples/examples/gaussian_processes/GP-Circular.html).
#
# ### Circular Kernel
#
# When the underlying space is polar, typical Euclidean kernels such as Matern kernels will be insufficient as the boundary. Circular kernels do not exhibit this behaviour and instead _wrap_ around the boundary points to create a smooth function. Such a kernel was given in [Padonou & Roustant (2015)](https://hal.inria.fr/hal-01119942v1) and for two angles $\theta, \theta'$ can be written as
# $$W_c(\theta, \theta') = \left\lvert \left(1 + \tau \frac{d(\theta, \theta')}{c} \right) \left(1 - \frac{d(\theta, \theta')}{c} \right)^{\tau} \right\rvert \quad \tau \geq 4 \tag{1}.$$
#
# Here we have the parameter $\tau$ which is analogous to a lengthscale for Euclidean stationary kernels as it controls the correlation between pairs of observations. Further, the distance metric $d$ here is an angular distance
# $$d(\theta, \theta') = \lvert (\theta-\theta'+c) \operatorname{mod} 2c - c \rvert.$$
#
# To implement this, one must write the following class 

# %%
from chex import dataclass


def angular_distance(x, y, c):
    return jnp.abs((x - y + c) % (c * 2) - c)
    

@dataclass(repr=False)
class Polar(gpx.kernels.Kernel):
    period: float = 2*jnp.pi

    def __post_init__(self):
        self.c = self.period / 2.0  # in [0, \pi]


    def __call__(self, x: gpx.types.Array, y: gpx.types.Array, params: dict) -> gpx.types.Array:
        tau = params["tau"]
        t = angular_distance(x, y, self.c)
        K = (1 + tau * t / self.c) * jnp.clip(1 - t / self.c, 0, jnp.inf) ** tau
        return K.squeeze()

    @property
    def params(self) -> dict:
        return {'tau': jnp.array([4.])}

# %% [markdown]
# We can unpack this now to make better sense of it. In the kernel's `__init__` function we are simply specifying the length of a single period. Here, the underlying domain is a circle, so one period is simply $2\pi$. Next we define the kernel's `__call__` function which is a direct implementation of Equation 1. Finally, we define the Kernel's parameter property which contains just one value $\tau$. Now $\tau$ is defined to be greater than or equal to 4, so we initialise it at 4. 
#
#
# ### Custom Parameter Bijection
#
# The constraint on $\tau$ will make optimisation challenging and it'd be much more helpful if we could transform $\tau$ to be on the full real line when we perform gradient descent on it. Fortunately, the `add_parameter` function in GPJax means we must simply define the parameter's name and corresponding bijection and everything else will be taken care of. Under the hood, a master configuration object is being updated when we call this function to register this parameter and its corresponding transform.

# %%
from gpjax.config import add_parameter
import tensorflow_probability.substrates.jax.bijectors as tfb

add_parameter('tau',  tfb.Softplus(low=jnp.array(4.0)))

non_informative
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
non_informative
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
slice_kernel = gpx.RBF(active_dims = [0, 1, 3])

# %% [markdown]
# This will implicitly make the kernel have a single lengthscale per input dimensions. Our kernel is now an ARD kernel:

# %%
print(f"ARD: {slice_kernel.ard}")
print(f"Lengthscales: {slice_kernel.params['lengthscale']}")

# %% [markdown]
# We'll now simulate some data and evaluate the kernel on the previously selected input dimensions

# %%
x_matrix = jr.normal(key, shape = (50, 5))
K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params)
print(K.shape)

# %% [markdown]
# ## Combination of Kernels
#
# The product or sum of two positive define matrices yields a positive definite matrix. Consequently, summing or multiplying sets of kernels is a valid operation and can yield rich kernel functions. In GPJax, sums of kernels can be created by passing a list of kernels into the `SumKernel` object in the following way.

# %%
k1 = gpx.RBF()
k1._params = {'lengthscale': jnp.array(1.), 'variance': jnp.array(10.)}
k2 = gpx.Polynomial()
sum_k = gpx.SumKernel(kernel_set = [k1, k2])

fig, ax = plt.subplots(ncols=3, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(sum_k, x, sum_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])

# %% [markdown]
# Similarily, products of kernels can be created through the `ProductKernel` object.

# %%
k3 = gpx.Matern32()
prod_k = gpx.ProductKernel(kernel_set = [k1, k2, k3])
fig, ax = plt.subplots(ncols=4, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(k3, x, k3.params))
im3=ax[3].matshow(gpx.kernels.gram(prod_k, x, prod_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])
fig.colorbar(im3, ax=ax[3])

# %% [markdown]
# ## Custom Kernel
#
# If one wishes to implement their own kernel then the process is straightforward. There are two steps to achieving this 
# 1. List the kernel's parameters
# 2. Define the kernel's pairwise operation
#
# We'll demonstrate this process now for a ciruclar kernel. This example is an adaption of the excellent guide given in the PYMC3 documentation and we would encourage the curious reader to visit their notebook which can be found [here](https://docs.pymc.io/pymc-examples/examples/gaussian_processes/GP-Circular.html).
#
# ### Circular Kernel
#
# When the underlying space is polar, typical Euclidean kernels such as Matern kernels will be insufficient as the boundary. Circular kernels do not exhibit this behaviour and instead _wrap_ around the boundary points to create a smooth function. Such a kernel was given in [Padonou & Roustant (2015)](https://hal.inria.fr/hal-01119942v1) and for two angles $\theta, \theta'$ can be written as
# $$W_c(\theta, \theta') = \left\lvert \left(1 + \tau \frac{d(\theta, \theta')}{c} \right) \left(1 - \frac{d(\theta, \theta')}{c} \right)^{\tau} \right\rvert \quad \tau \geq 4 \tag{1}.$$
#
# Here we have the parameter $\tau$ which is analogous to a lengthscale for Euclidean stationary kernels as it controls the correlation between pairs of observations. Further, the distance metric $d$ here is an angular distance
# $$d(\theta, \theta') = \lvert (\theta-\theta'+c) \operatorname{mod} 2c - c \rvert.$$
#
# To implement this, one must write the following class 

# %%
from chex import dataclass


def angular_distance(x, y, c):
    return jnp.abs((x - y + c) % (c * 2) - c)
    

@dataclass(repr=False)
class Polar(gpx.kernels.Kernel):
    period: float = 2*jnp.pi

    def __post_init__(self):
        self.c = self.period / 2.0  # in [0, \pi]


    def __call__(self, x: gpx.types.Array, y: gpx.types.Array, params: dict) -> gpx.types.Array:
        tau = params["tau"]
        t = angular_distance(x, y, self.c)
        K = (1 + tau * t / self.c) * jnp.clip(1 - t / self.c, 0, jnp.inf) ** tau
        return K.squeeze()

random
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
slice_kernel = gpx.RBF(active_dims = [0, 1, 3])

# %% [markdown]
# This will implicitly make the kernel have a single lengthscale per input dimensions. Our kernel is now an ARD kernel:

# %%
print(f"ARD: {slice_kernel.ard}")
print(f"Lengthscales: {slice_kernel.params['lengthscale']}")

# %% [markdown]
# We'll now simulate some data and evaluate the kernel on the previously selected input dimensions

# %%
x_matrix = jr.normal(key, shape = (50, 5))
K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params)
print(K.shape)

# %% [markdown]
# ## Combination of Kernels
#
# The product or sum of two positive define matrices yields a positive definite matrix. Consequently, summing or multiplying sets of kernels is a valid operation and can yield rich kernel functions. In GPJax, sums of kernels can be created by passing a list of kernels into the `SumKernel` object in the following way.

# %%
k1 = gpx.RBF()
k1._params = {'lengthscale': jnp.array(1.), 'variance': jnp.array(10.)}
k2 = gpx.Polynomial()
sum_k = gpx.SumKernel(kernel_set = [k1, k2])

fig, ax = plt.subplots(ncols=3, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(sum_k, x, sum_k.params))


fig.colorbar(im0, ax=ax[0])
random
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


random
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
slice_kernel = gpx.RBF(active_dims = [0, 1, 3])

# %% [markdown]
# This will implicitly make the kernel have a single lengthscale per input dimensions. Our kernel is now an ARD kernel:

# %%
print(f"ARD: {slice_kernel.ard}")
print(f"Lengthscales: {slice_kernel.params['lengthscale']}")

# %% [markdown]
# We'll now simulate some data and evaluate the kernel on the previously selected input dimensions

# %%
x_matrix = jr.normal(key, shape = (50, 5))
K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params)
print(K.shape)

# %% [markdown]
# ## Combination of Kernels
#
# The product or sum of two positive define matrices yields a positive definite matrix. Consequently, summing or multiplying sets of kernels is a valid operation and can yield rich kernel functions. In GPJax, sums of kernels can be created by passing a list of kernels into the `SumKernel` object in the following way.

# %%
k1 = gpx.RBF()
k1._params = {'lengthscale': jnp.array(1.), 'variance': jnp.array(10.)}
k2 = gpx.Polynomial()
sum_k = gpx.SumKernel(kernel_set = [k1, k2])

fig, ax = plt.subplots(ncols=3, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(sum_k, x, sum_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])

# %% [markdown]
# Similarily, products of kernels can be created through the `ProductKernel` object.

# %%
k3 = gpx.Matern32()
prod_k = gpx.ProductKernel(kernel_set = [k1, k2, k3])
fig, ax = plt.subplots(ncols=4, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(k3, x, k3.params))
im3=ax[3].matshow(gpx.kernels.gram(prod_k, x, prod_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])
random
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.11.5
#   kernelspec:
#     display_name: Python 3.9.7 ('gpjax')
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Guide to Kernels
#
# In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.

# %%
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from optax import adam
from jax import jit 

key = jr.PRNGKey(123)

# %% [markdown]
# ## Supported Kernels 
#
# The following kernels are supported in GPJax
# * Matern - first, second and third order
# * RBF or Squared Exponential
# * Polynomial
#
# The syntax for using these kernels is consistent across objects. However, the kernel's type influences the type of sample paths that can be drawn. To visualise this, we plot a set of 10 function draws per kernel in the following cell.

# %%
kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]
fig, axes = plt.subplots(ncols=3, nrows=2, figsize=(20, 10))

x = jnp.linspace(-3., 3., num=200).reshape(-1, 1)


for k, ax in zip(kernels, axes.ravel()):
    prior = gpx.Prior(kernel=k)
    params, _, _ = gpx.initialise(prior)
    rv = prior.random_variable(x, params)
    y = rv.sample(sample_shape = 10, seed = key)

    ax.plot(x, y.T, alpha=0.5)
    ax.set_title(k.name)

# %% [markdown]
# ### Active dimensions
#
# By default, kernels will operate over every dimension of the supplied inputs. For some use cases it is desirable to have the kernel operate over specific dimensions of the input data. Every kernel has an `active_dims` attribute that specificies the index values that the kernel should evaluate. 
#
# To see this, consider the following 5-dimensional dataset for which we would like our RBF kernel to act on the first, second and fourth dimensions.

# %%
slice_kernel = gpx.RBF(active_dims = [0, 1, 3])

# %% [markdown]
# This will implicitly make the kernel have a single lengthscale per input dimensions. Our kernel is now an ARD kernel:

# %%
print(f"ARD: {slice_kernel.ard}")
print(f"Lengthscales: {slice_kernel.params['lengthscale']}")

# %% [markdown]
# We'll now simulate some data and evaluate the kernel on the previously selected input dimensions

# %%
x_matrix = jr.normal(key, shape = (50, 5))
K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params)
print(K.shape)

# %% [markdown]
# ## Combination of Kernels
#
# The product or sum of two positive define matrices yields a positive definite matrix. Consequently, summing or multiplying sets of kernels is a valid operation and can yield rich kernel functions. In GPJax, sums of kernels can be created by passing a list of kernels into the `SumKernel` object in the following way.

# %%
k1 = gpx.RBF()
k1._params = {'lengthscale': jnp.array(1.), 'variance': jnp.array(10.)}
k2 = gpx.Polynomial()
sum_k = gpx.SumKernel(kernel_set = [k1, k2])

fig, ax = plt.subplots(ncols=3, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(sum_k, x, sum_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])

# %% [markdown]
# Similarily, products of kernels can be created through the `ProductKernel` object.

# %%
k3 = gpx.Matern32()
prod_k = gpx.ProductKernel(kernel_set = [k1, k2, k3])
fig, ax = plt.subplots(ncols=4, figsize=(20, 5))
im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params))
im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params))
im2=ax[2].matshow(gpx.kernels.gram(k3, x, k3.params))
im3=ax[3].matshow(gpx.kernels.gram(prod_k, x, prod_k.params))


fig.colorbar(im0, ax=ax[0])
fig.colorbar(im1, ax=ax[1])
fig.colorbar(im2, ax=ax[2])
fig.colorbar(im3, ax=ax[3])

# %% [markdown]
# ## Custom Kernel
#
# If one wishes to implement their own kernel then the process is straightforward. There are two steps to achieving this 
# 1. List the kernel's parameters
# 2. Define the kernel's pairwise operation
#
# We'll demonstrate this process now for a ciruclar kernel. This example is an adaption of the excellent guide given in the PYMC3 documentation and we would encourage the curious reader to visit their notebook which can be found [here](https://docs.pymc.io/pymc-examples/examples/gaussian_processes/GP-Circular.html).
#
# ### Circular Kernel
#
# When the underlying space is polar, typical Euclidean kernels such as Matern kernels will be insufficient as the boundary. Circular kernels do not exhibit this behaviour and instead _wrap_ around the boundary points to create a smooth function. Such a kernel was given in [Padonou & Roustant (2015)](https://hal.inria.fr/hal-01119942v1) and for two angles $\theta, \theta'$ can be written as
# $$W_c(\theta, \theta') = \left\lvert \left(1 + \tau \frac{d(\theta, \theta')}{c} \right) \left(1 - \frac{d(\theta, \theta')}{c} \right)^{\tau} \right\rvert \quad \tau \geq 4 \tag{1}.$$
#
# Here we have the parameter $\tau$ which is analogous to a lengthscale for Euclidean stationary kernels as it controls the correlation between pairs of observations. Further, the distance metric $d$ here is an angular distance
# $$d(\theta, \theta') = \lvert (\theta-\theta'+c) \operatorname{mod} 2c - c \rvert.$$
#
# To implement this, one must write the following class 

# %%
from chex import dataclass


def angular_distance(x, y, c):
    return jnp.abs((x - y + c) % (c * 2) - c)
    

@dataclass(repr=False)
class Polar(gpx.kernels.Kernel):
    period: float = 2*jnp.pi

    def __post_init__(self):
        self.c = self.period / 2.0  # in [0, \pi]


    def __call__(self, x: gpx.types.Array, y: gpx.types.Array, params: dict) -> gpx.types.Array:
        tau = params["tau"]
        t = angular_distance(x, y, self.c)
        K = (1 + tau * t / self.c) * jnp.clip(1 - t / self.c, 0, jnp.inf) ** tau
        return K.squeeze()

    @property
    def params(self) -> dict:
        return {'tau': jnp.array([4.])}

# %% [markdown]
# We can unpack this now to make better sense of it. In the kernel's `__init__` function we are simply specifying the length of a single period. Here, the underlying domain is a circle, so one period is simply $2\pi$. Next we define the kernel's `__call__` function which is a direct implementation of Equation 1. Finally, we define the Kernel's parameter property which contains just one value $\tau$. Now $\tau$ is defined to be greater than or equal to 4, so we initialise it at 4. 
#
#
# ### Custom Parameter Bijection
#
# The constraint on $\tau$ will make optimisation challenging and it'd be much more helpful if we could transform $\tau$ to be on the full real line when we perform gradient descent on it. Fortunately, the `add_parameter` function in GPJax means we must simply define the parameter's name and corresponding bijection and everything else will be taken care of. Under the hood, a master configuration object is being updated when we call this function to register this parameter and its corresponding transform.

# %%
from gpjax.config import add_parameter
import tensorflow_probability.substrates.jax.bijectors as tfb

add_parameter('tau',  tfb.Softplus(low=jnp.array(4.0)))

# %% [markdown]
# ### Using our Polar Kernel
#
# With a custom kernel now defined, we can go ahead and fit our new kernel. To do this, we randomly simulate a sequence of points on a circle and fit our GP to this using the above defined circular kernel. For further details on this process, it's advised to see the [Regression notebook](https://gpjax.readthedocs.io/en/latest/nbs/regression.html).

# %%
# Simulate data
angles = jnp.linspace(0, 2 * jnp.pi, num=200).reshape(-1, 1)
n = 20
noise = 0.2

X = jnp.sort(jr.uniform(key, minval=0.0, maxval=jnp.pi * 2, shape=(n, 1)), axis=0)
y = 4 + jnp.cos(2 * X) + jr.normal(key, shape=X.shape) * noise
training = gpx.Dataset(X=X, y=y)

# Define Polar Gaussian process
PKern = Polar()
circlular_posterior = gpx.Prior(kernel=PKern) * gpx.Gaussian(num_datapoints=n)

# Initialise parameters and corresponding transformations
params, constrainer, unconstrainer = gpx.initialise(circlular_posterior)

# Optimise GP's marginal log-likelihood using Adam
mll = jit(circlular_posterior.marginal_log_likelihood(training, constrainer, negative=True))
learned_params = gpx.optax_fit(mll, params, adam(learning_rate=0.05), n_iters=1000)

# Untransform learned parameters
final_params = gpx.transform(learned_params, constrainer)

# %% [markdown]
# ### Prediction
#
# We'll now query the GP's predictive posterior distribution at a linearly spaced set of test points and plot the results.

# %%
mu = circlular_posterior.mean(training, final_params)(angles).squeeze()
one_sigma = jnp.sqrt(jnp.diag(circlular_posterior.variance(training, final_params)(angles)))

# %%
fig = plt.figure(figsize=(10, 8))
gridspec = fig.add_gridspec(1, 1)
ax = plt.subplot(gridspec[0], polar=True)

ax.grid(color="#888888")  # Color the grid
ax.spines["polar"].set_visible(False)  # Show or hide the plot spine
ax.fill_between(
    angles.squeeze(),
    mu - one_sigma,
    mu + one_sigma,
    alpha=0.3,
    label=r"1 Posterior s.d.",
    color="#B5121B"
)
random
PKern = Polar()         tau = params["tau"]         t = angular_distance(x, y, self.c) kernels = [gpx.Matern12(), gpx.Matern32(),gpx.Matern52(), gpx.RBF(), gpx.Polynomial(degree=1), gpx.Polynomial(degree=2)]     prior = gpx.Prior(kernel=k)     params, _, _ = gpx.initialise(prior)     rv = prior.random_variable(x, params) add_parameter('tau',  tfb.Softplus(low=jnp.array(4.0)))     ax.plot(x, y.T, alpha=0.5) n = 20 X = jnp.sort(jr.uniform(key, minval=0.0, maxval=jnp.pi * 2, shape=(n, 1)), axis=0) training = gpx.Dataset(X=X, y=y) slice_kernel = gpx.RBF(active_dims = [0, 1, 3]) circlular_posterior = gpx.Prior(kernel=PKern) * gpx.Gaussian(num_datapoints=n) params, constrainer, unconstrainer = gpx.initialise(circlular_posterior) mll = jit(circlular_posterior.marginal_log_likelihood(training, constrainer, negative=True)) learned_params = gpx.optax_fit(mll, params, adam(learning_rate=0.05), n_iters=1000) K = gpx.kernels.gram(slice_kernel, x_matrix, slice_kernel.params) final_params = gpx.transform(learned_params, constrainer) mu = circlular_posterior.mean(training, final_params)(angles).squeeze() k1 = gpx.RBF() one_sigma = jnp.sqrt(jnp.diag(circlular_posterior.variance(training, final_params)(angles))) k2 = gpx.Polynomial() sum_k = gpx.SumKernel(kernel_set = [k1, k2]) im0=ax[0].matshow(gpx.kernels.gram(k1, x, k1.params)) im1=ax[1].matshow(gpx.kernels.gram(k2, x, k2.params)) im2=ax[2].matshow(gpx.kernels.gram(sum_k, x, sum_k.params)) k3 = gpx.Matern32() prod_k = gpx.ProductKernel(kernel_set = [k1, k2, k3]) im2=ax[2].matshow(gpx.kernels.gram(k3, x, k3.params)) im3=ax[3].matshow(gpx.kernels.gram(prod_k, x, prod_k.params)) ax.plot(angles, mu, label='Posterior mean') # The following kernels are supported in GPJax # %% [markdown] # In this notebook we'll expose the kernels available in GPJax and show how a custom polar kernel can be implemented.     @property fig.colorbar(im1, ax=ax[1]) for k, ax in zip(kernels, axes.ravel()): fig.colorbar(im3, ax=ax[3]) ax.fill_between(