Center-gap GPQR

In this example, the data generating process is $$Y = \cos(2 \pi (X+0.1)) + \epsilon, \quad \epsilon \sim \mathcal{N}(0, X+0.1),$$ and prior mean is $$ \mu(X; \theta) = \cos(2 \pi X + \theta) $$ therefore residual is $$ R = Y - \mu(X; \theta). $$

Central quantile $Q_{\tau_0}(x)$ and gaps $\Delta Q_{\tau_i}(x)$ of $R$ are modeled by linear combination of latent GP $g_j(x)$:

$$Q_{\tau_0}(x) = \sum_j a_{0j}g_j(x), \quad \Delta Q_{\tau_i}(x) = \log \left(1 + \exp \sum_j a_{ij}g_j(x)\right),$$ where $$ g_j(x) \sim \mathcal{N}(c_j, k(x, x’)). $$

Parameters $\theta$, $a_{ij}$ and $c_j$ are learned by maximizing the marginal likelihood.

import os

import torch
from torch.distributions import Normal
from gpytorch.variational import CholeskyVariationalDistribution
from gpytorch.variational import VariationalStrategy
from gpytorch.variational import LMCVariationalStrategy
from gpytorch.means import ConstantMean
from gpytorch.kernels import RBFKernel, ScaleKernel
from gpytorch.mlls import VariationalELBO
import matplotlib.pyplot as plt

from gpytorch_qr.models import CenterGapQuantileGP
from gpytorch_qr.likelihoods import CenterGapQuantileLikelihood

try:
    import sys

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

    import config_notebook
except ImportError:
    print("Output will not be deterministic SVG.")

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

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

Data preparation

def f(x):
    return torch.cos((x + 0.1) * 2 * 3.14)


def std(x):
    return x + 0.1


x_range = torch.linspace(0, 1, 100).reshape(-1, 1).to(device)
x = x_range.repeat(5, 1)
y = (f(x) + torch.randn(x.shape, device=device).mul(std(x))).squeeze()

q = torch.tensor([0.1, 0.25, 0.5, 0.75, 0.9]).to(device)
true_quantiles = f(x_range) + std(x_range) * Normal(0, 1).icdf(q)

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

class PriorMean(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.theta = torch.nn.Parameter(torch.tensor(0.0))

    def forward(self, x):
        return torch.cos(2 * 3.14 * x + self.theta)


mean = PriorMean().to(device)

with torch.no_grad():
    res = y - mean(x).squeeze(-1)
    m = mean(x_range)
    res_quantiles = true_quantiles - m

fig, axes = plt.subplots(1, 2, figsize=(12, 4))

axes[0].scatter(x.cpu(), y.cpu(), c="k", marker=".")
axes[0].plot(x_range.cpu(), true_quantiles.cpu(), "--", c="gray")
axes[0].plot(x_range.cpu(), m.detach().cpu())

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

plt.show()

Define models and likelihoods

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

        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_q_index = (q - 0.5).abs().argmin().item()
num_latents = len(q) - 2
gp = MyGP(inducing_points, len(q), central_q_index, num_latents).to(device)
likelihood = CenterGapQuantileLikelihood(q, central_q_index).to(device)

Train

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

for _ in range(n_epochs):
    optimizer.zero_grad()
    res = y - mean(x).squeeze(-1)
    output = gp(x)
    loss = -mll(output, res)
    loss.backward()
    optimizer.step()

Evaluate

mean.eval()
gp.eval()
with torch.no_grad():
    m = mean(x_pred)
    mean_q = gp.mean_quantiles_mc(x_pred) + m
    lower_q, upper_q = (
        gp.quantile_quantiles_mc(x_pred, torch.tensor([0.025, 0.975]).to(device)) + m
    )

Plot result

colors = plt.cm.tab10.colors

plt.scatter(x.cpu(), y.cpu(), c="gray", marker=".", alpha=0.1)
plt.plot(x_range.cpu(), true_quantiles.cpu(), "--", c="k")

for i in range(len(q)):
    plt.plot(x_pred.cpu(), mean_q[:, i].cpu(), color=colors[i])
    plt.fill_between(
        x_pred.cpu().squeeze(),
        lower_q[:, i].cpu(),
        upper_q[:, i].cpu(),
        color=colors[i],
        alpha=0.3,
    )
plt.show()