MultivariateGPR interface prototype

[tadejkrivec]

PR type: new feature

Related issue(s)/PRs: #1555, #1209, #1175

Summary

Proposed changes

Implementation of MultivariateGPR interface for multi-output regression with GPR model.

• MultivariateGPR prediction is based on the "inducing_point_conditional" implementation which can handle different covariance shapes based on the input arguments (full_cov, full_output_cov).
• Added the jitter on the diagonal of the training covariance matrix.
• Tests are not implemented

What alternatives have you considered?

No alternatives were considered, but I am wondering if the code in "inducing_point_conditional" could be somehow reused instead of rewriten in MultivariateGPR prediction.

Minimal working example

import tensorflow as tf
import numpy as np
import gpflow
from gpflow.models import training_loss_closure
from gpflow.ci_utils import ci_niter
import matplotlib.pyplot as plt

def optimize_model_with_scipy(model):
    opt = gpflow.optimizers.Scipy()
    loss_closure = training_loss_closure(model, data, compile=True)   
    _ = opt.minimize(
        loss_closure,
        variables=model.trainable_variables,
        method="l-bfgs-b",
        options=dict(maxiter=ci_niter(10000)),
        compile=True
    )

def plot_model(model, lower=-8.0, upper=8.0, n=100):
    pX = np.linspace(lower, upper, n)[:, None]
    pY, pYv = model.predict_f(pX)
    pY, pYv = pY.numpy(), pYv.numpy()
    plt.plot(X, Y, "kx")
    plt.gca().set_prop_cycle(None)
    plt.plot(pX, pY, 'rx')
    for i in range(pY.shape[1]):
        top = pY[:, i] + 2.0 * pYv[:, i] ** 0.5
        bot = pY[:, i] - 2.0 * pYv[:, i] ** 0.5
        plt.fill_between(pX[:, 0], top, bot, alpha=0.3)
    plt.xlabel("X")
    plt.ylabel("f")
    
def generate_data(N=100):
    X = np.random.rand(N)[:, None] * 10 - 5  # Inputs = N x D
    G = np.hstack((0.5 * np.sin(3 * X) + X, 3.0 * np.cos(X) - X))  # G = N x L
    W = np.array([[0.5, -0.3, 1.5], [-0.4, 0.43, 0.0]])  # L x P
    F = np.matmul(G, W)  # N x P
    Y = F + np.random.randn(*F.shape) * [0.2, 0.2, 0.2]
    return X, Y

N = 100  # number of points
D = 1  # number of input dimensions
M = 15  # number of inducing points
L = 2  # number of latent GPs
P = 3  # number of observations = output dimensions
X, Y = data = generate_data(N)
Zinit = X.copy()

kern = gpflow.kernels.RBF()
kernel = gpflow.kernels.SharedIndependent(kern, output_dim = Y.shape[1])
model = gpflow.models.MultivariateGPR(kernel = kernel, data = data)
optimize_model_with_scipy(model)
pY, pYv = model.predict_f(X[:80, :], full_cov = True, full_output_cov = True)
print('full_cov = True,', 'full_output_cov = True', pY.shape, pYv.shape)
pY, pYv = model.predict_f(X[:80, :], full_cov = True, full_output_cov = False)
print('full_cov = True,', 'full_output_cov = False', pY.shape, pYv.shape)
pY, pYv = model.predict_f(X[:80, :], full_cov = False, full_output_cov = True)
print('full_cov = False,', 'full_output_cov = True', pY.shape, pYv.shape)
pY, pYv = model.predict_f(X[:80, :], full_cov = False, full_output_cov = False)
print('full_cov = False,', 'full_output_cov = False', pY.shape, pYv.shape)
plot_model(model)

kern_list = [gpflow.kernels.RBF() for _ in range(Y.shape[1])]
kernel = gpflow.kernels.SeparateIndependent(kern_list)
model = gpflow.models.MultivariateGPR(kernel=kernel, data = data)
optimize_model_with_scipy(model)
pY, pYv = model.predict_f(X[:80, :], full_cov = True, full_output_cov = True)
print('full_cov = True,', 'full_output_cov = True', pY.shape, pYv.shape)
pY, pYv = model.predict_f(X[:80, :], full_cov = True, full_output_cov = False)
print('full_cov = True,', 'full_output_cov = False', pY.shape, pYv.shape)
pY, pYv = model.predict_f(X[:80, :], full_cov = False, full_output_cov = True)
print('full_cov = False,', 'full_output_cov = True', pY.shape, pYv.shape)
pY, pYv = model.predict_f(X[:80, :], full_cov = False, full_output_cov = False)
print('full_cov = False,', 'full_output_cov = False', pY.shape, pYv.shape)
plot_model(model)

L = 2 
kern_list = [gpflow.kernels.RBF() for _ in range(L)]
kernel = gpflow.kernels.LinearCoregionalization(
    kern_list, W=np.random.randn(P, L)
)
model = gpflow.models.MultivariateGPR(kernel=kernel, data = data)
optimize_model_with_scipy(model)
pY, pYv = model.predict_f(X[:80, :], full_cov = True, full_output_cov = True)
print('full_cov = True,', 'full_output_cov = True', pY.shape, pYv.shape)
pY, pYv = model.predict_f(X[:80, :], full_cov = True, full_output_cov = False)
print('full_cov = True,', 'full_output_cov = False', pY.shape, pYv.shape)
pY, pYv = model.predict_f(X[:80, :], full_cov = False, full_output_cov = True)
print('full_cov = False,', 'full_output_cov = True', pY.shape, pYv.shape)
pY, pYv = model.predict_f(X[:80, :], full_cov = False, full_output_cov = False)
print('full_cov = False,', 'full_output_cov = False', pY.shape, pYv.shape)
plot_model(model)

PR checklist

• New features: code is well-documented
  • detailed docstrings (API documentation)
  • notebook examples (usage demonstration)
• The bug case / new feature is covered by unit tests
• Code has type annotations
• I ran the black+isort formatter (make format)
• I locally tested that the tests pass (make check-all)

Release notes

Fully backwards compatible: yes