Multi-output GPQR

import os
import torch
import matplotlib.pyplot as plt

torch.manual_seed(42)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

n_epochs = int(os.getenv("GPYTORCHQR_N_EPOCHS", 5000))

Input data

def mean1(x):
    return torch.cos(x.squeeze(-1) * 2 * 3.14)


def mean2(x):
    return torch.sin(x.squeeze(-1) * 2 * 3.14)


def std(x):
    return x + 0.1


x_range = torch.linspace(0, 1, 100, device=device).reshape(-1, 1)
x = x_range.repeat(2, 1)
y1 = mean1(x).unsqueeze(-1) + torch.randn(x.shape, device=device).mul(std(x))
y2 = mean2(x).unsqueeze(-1) + torch.randn(x.shape, device=device).mul(std(x))
y = torch.concatenate([y1, y2], dim=-1)

q1 = torch.tensor([0.1, 0.5, 0.9], device=device)
icdf1 = torch.distributions.Normal(0, std(x_range)).icdf(q1)
true_quantiles1 = mean1(x_range).unsqueeze(-1) + icdf1

q2 = torch.tensor([0.1, 0.25, 0.5, 0.75, 0.9], device=device)
icdf2 = torch.distributions.Normal(0, std(x_range)).icdf(q2)
true_quantiles2 = mean2(x_range).unsqueeze(-1) + icdf2

x_pred = torch.linspace(0, 1.5, 100).reshape(-1, 1).to(device)

fig, axes = plt.subplots(1, 2)

axes[0].scatter(x.cpu(), y1.cpu(), c="k", marker=".")
axes[0].plot(x_range.cpu(), true_quantiles1.cpu(), "--", c="gray")

axes[1].scatter(x.cpu(), y2.cpu(), c="k", marker=".")
axes[1].plot(x_range.cpu(), true_quantiles2.cpu(), "--", c="gray")

fig.show()

Multitask model

from gpytorch.means import ConstantMean
from gpytorch.kernels import ScaleKernel, RBFKernel
from gpytorch.variational import (
    CholeskyVariationalDistribution,
    VariationalStrategy,
)
from gpytorch_qr.models import CenterGapQuantileGP
from gpytorch_qr.likelihoods import MultioutputCenterGapQuantilesLikelihood
from gpytorch_qr.variational import CenterGapLMCVariationalStrategy


class QuantileGPModel(CenterGapQuantileGP):
    def __init__(
        self,
        inducing_points,
        num_quantiles,
        num_lower_quantiles,
        num_latents,
        num_central_latents,
    ):
        N, D = inducing_points.size()
        variational_distribution = CholeskyVariationalDistribution(
            N,
            batch_shape=torch.Size([num_latents]),
        )
        variational_strategy = CenterGapLMCVariationalStrategy(
            VariationalStrategy(
                self,
                inducing_points,
                variational_distribution,
                learn_inducing_locations=True,
            ),
            sum(num_quantiles),
            num_latents,
            num_central_latents=num_central_latents,
            num_quantiles=num_quantiles,
        )

        mean = ConstantMean(batch_shape=torch.Size([num_latents]))
        covar = ScaleKernel(
            RBFKernel(ard_num_dims=D, batch_shape=torch.Size([num_latents])),
            batch_shape=torch.Size([num_latents]),
        )
        super().__init__(
            variational_strategy, mean, covar, num_quantiles, num_lower_quantiles
        )


inducing_points = torch.linspace(0, 1, 10).reshape(-1, 1).to(device)
central_q1_index = (q1 - 0.5).abs().argmin().item()
central_q2_index = (q2 - 0.5).abs().argmin().item()
num_latents = 7
num_central_latents = 3

likelihood = MultioutputCenterGapQuantilesLikelihood(
    [q1, q2],
    [central_q1_index, central_q2_index],
).to(device)
model = QuantileGPModel(
    inducing_points,
    [len(q1), len(q2)],
    [central_q1_index, central_q2_index],
    num_latents,
    num_central_latents,
).to(device)

from gpytorch.mlls import VariationalELBO

model.train()
likelihood.train()

optimizer = torch.optim.Adam(
    list(model.parameters()) + list(likelihood.parameters()),
    lr=0.001,
)
mll = VariationalELBO(likelihood, model, num_data=y.numel())

for _ in range(n_epochs):
    output = model(x)
    loss = -mll(output, y)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

model.eval()
likelihood.eval()

with torch.no_grad():
    mean_q = model.mean_quantiles_mc(x_pred)
    lower_q, upper_q = model.quantile_quantiles_mc(
        x_pred, torch.tensor([0.025, 0.975], device=device)
    )

pred_mean_q = mean_q.detach().cpu()
pred_lower_q = lower_q.detach().cpu()
pred_upper_q = upper_q.detach().cpu()

fig, axes = plt.subplots(1, 2)

axes[0].scatter(x.cpu(), y1.cpu(), c="gray")
axes[0].plot(x_range.cpu(), true_quantiles1.cpu(), color="k")
for i in range(len(q1)):
    axes[0].plot(x_pred.cpu(), mean_q[:, i].cpu(), label=f"q={q1[i].item():.2f}")
    axes[0].fill_between(
        x_pred.cpu().squeeze(),
        lower_q[:, i].cpu(),
        upper_q[:, i].cpu(),
        alpha=0.3,
    )

axes[1].scatter(x.cpu(), y2.cpu(), c="gray")
axes[1].plot(x_range.cpu(), true_quantiles2.cpu(), color="k")
for i in range(len(q2)):
    axes[1].plot(
        x_pred.cpu(), mean_q[:, len(q1) + i].cpu(), label=f"q={q2[i].item():.2f}"
    )
    axes[1].fill_between(
        x_pred.cpu().squeeze(),
        lower_q[:, len(q1) + i].cpu(),
        upper_q[:, len(q1) + i].cpu(),
        alpha=0.3,
    )

fig.show()