<jupyter_start><jupyter_text>Spectral Kernel Approximations<jupyter_code>import gpjax
from gpjax.gps import Prior
from gpjax.likelihoods import Gaussian
import jax
import jax.numpy as jnp
import jax.random as jr
from jax.experimental import optimizers
import matplotlib.pyplot as plt

key = jr.PRNGKey(123)
def draw_gp_sample(key, N, n_train, kernel, params, noise=0.1):
    x = jr.uniform(key=key, minval=-3.0, maxval=3.0, shape=(N,)).sort().reshape(-1, 1)
    mean = jnp.zeros(N)
    cov = gpjax.kernels.gram(kernel, x, params) + gpjax.utils.I(N) * 1e-6
    y = jr.multivariate_normal(key, mean=mean, cov=cov)

    train_idx = jr.choice(
        key=key, a=jnp.arange(N), shape=(n_train,), replace=False
    ).sort()
    xtr = x[train_idx]
    ytr = y[train_idx]
    ytr += jr.normal(key, shape=ytr.shape) * noise
    return xtr, x, ytr, y
N = 2000
n_train = 1000
noise = 0.1
M = 100
x, xte, y, yte = draw_gp_sample(
    key,
    N,
    n_train,
    gpjax.kernels.RBF(),
    {"lengthscale": jnp.array(0.4), "variance": jnp.array(0.5)},
    noise,
)
fig, ax = plt.subplots(figsize=(10, 5))
ax.plot(x, y, "o", label="Observations", alpha=0.5)
ax.plot(xte, yte, label="Latent function")
ax.legend(loc="best")
# plt.savefig("data.png")
from gpjax.kernels import RBF, initialise
from gpjax.kernels.spectral import to_spectral

kernel = to_spectral(RBF(), M)
print(kernel)
from gpjax import Prior

prior = Prior(kernel=kernel)
from gpjax.likelihoods import Gaussian

posterior = prior * Gaussian()
print(posterior)
from gpjax.parameters import initialise, build_all_transforms
from gpjax.config import get_defaults

params = initialise(key, posterior)
sparams = {"basis_fns": params["basis_fns"]}
del params["basis_fns"]
constrain, unconstrain = build_all_transforms(params.keys(), get_defaults())
params = unconstrain(params)
from gpjax.objectives import marginal_ll
from jax import jit

mll = jit(marginal_ll(posterior, transform=constrain, negative=True))
mll(params, x, y, sparams)
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 = jax.value_and_grad(mll)(p, x, y, sparams)
    return opt_update(i, g, opt_state), v


nits = 300
mlls = []

for i in range(nits):
    opt_state, mll_estimate = step(i, opt_state)
    mlls.append(-mll_estimate)
fig, ax = plt.subplots(figsize=(10, 5))
ax.plot(range(1, nits + 1), mlls)
final_params = constrain(get_params(opt_state))
from gpjax.sampling import random_variable

posterior_rv = random_variable(
    posterior, final_params, x, y, xte, static_params=sparams
)
final_params
mu = posterior_rv.mean()
sigma = jnp.sqrt(jnp.diag(posterior_rv.covariance()))
fig, ax = plt.subplots(figsize=(10, 5))
ax.plot(x, y, "o", color="tab:orange")
ax.plot(xte, mu, color="tab:blue")
ax.fill_between(
    xte.ravel(),
    mu.ravel() - sigma,
    mu.ravel() + sigma,
    alpha=0.4,
    color="tab:blue",
)
ax.plot(xte, mu.ravel() - sigma, linestyle="--", color="tab:blue")
ax.plot(xte, mu.ravel() + sigma, linestyle="--", color="tab:blue")<jupyter_output><empty_output><jupyter_text>System information<jupyter_code>%load_ext watermark
%watermark -n -u -v -iv -w<jupyter_output>Last updated: Sun Mar 14 2021

Python implementation: CPython
Python version       : 3.8.7
IPython version      : 7.20.0

matplotlib: 3.3.4
jax       : 0.2.8
gpjax     : 0.3.3

Watermark: 2.2.0
{'fraction_non_alphanumeric': 0.1449530516431925, 'fraction_numerical': 0.022300469483568074, 'mean_word_length': 4.653399668325042, 'pattern_counts': {'":': 3, '<': 8, '<?xml version=': 0, '>': 8, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '69357', 'license': 'Apache-2.0', 'n_tokens_mistral': 1425, 'n_tokens_neox': 1339, 'n_words': 300, 'path': 'docs/nbs/spectral.ipynb', 'repo_name': 'jejjohnson/GPJax-1', 'size': '3408'}