{
"cells": [
{
"cell_type": "markdown",
"id": "0",
"metadata": {},
"source": [
"# Center-gap GPQR\n",
"\n",
"In this example, the data generating process is\n",
"$$Y = \\cos(2 \\pi (X+0.1)) + \\epsilon, \\quad \\epsilon \\sim \\mathcal{N}(0, X+0.1),$$\n",
"and prior mean is\n",
"$$ \\mu(X; \\theta) = \\cos(2 \\pi X + \\theta) $$\n",
"therefore residual is\n",
"$$ R = Y - \\mu(X; \\theta). $$\n",
"\n",
"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)$:\n",
"\n",
"$$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),$$\n",
"where \n",
"$$ g_j(x) \\sim \\mathcal{N}(c_j, k(x, x')). $$\n",
"\n",
"Parameters $\\theta$, $a_{ij}$ and $c_j$ are learned by maximizing the marginal likelihood."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1",
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"\n",
"import torch\n",
"from torch.distributions import Normal\n",
"from gpytorch.variational import CholeskyVariationalDistribution\n",
"from gpytorch.variational import VariationalStrategy\n",
"from gpytorch.variational import LMCVariationalStrategy\n",
"from gpytorch.means import ConstantMean\n",
"from gpytorch.kernels import RBFKernel, ScaleKernel\n",
"from gpytorch.mlls import VariationalELBO\n",
"import matplotlib.pyplot as plt\n",
"\n",
"from gpytorch_qr.models import CenterGapQuantileGP\n",
"from gpytorch_qr.likelihoods import CenterGapQuantileLikelihood\n",
"\n",
"try:\n",
" import sys\n",
"\n",
" sys.path.insert(0, os.path.abspath(\"../..\"))\n",
"\n",
" import config_notebook\n",
"except ImportError:\n",
" print(\"Output will not be deterministic SVG.\")\n",
"\n",
"torch.manual_seed(42)\n",
"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
"\n",
"n_epochs = int(os.getenv(\"GPYTORCHQR_N_EPOCHS\", 10000))"
]
},
{
"cell_type": "markdown",
"id": "2",
"metadata": {},
"source": [
"## Data preparation"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3",
"metadata": {},
"outputs": [],
"source": [
"def f(x):\n",
" return torch.cos((x + 0.1) * 2 * 3.14)\n",
"\n",
"\n",
"def std(x):\n",
" return x + 0.1\n",
"\n",
"\n",
"x_range = torch.linspace(0, 1, 100).reshape(-1, 1).to(device)\n",
"x = x_range.repeat(5, 1)\n",
"y = (f(x) + torch.randn(x.shape, device=device).mul(std(x))).squeeze()\n",
"\n",
"q = torch.tensor([0.1, 0.25, 0.5, 0.75, 0.9]).to(device)\n",
"true_quantiles = f(x_range) + std(x_range) * Normal(0, 1).icdf(q)\n",
"\n",
"x_pred = torch.linspace(0, 1.5, 100).reshape(-1, 1).to(device)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4",
"metadata": {},
"outputs": [],
"source": [
"class PriorMean(torch.nn.Module):\n",
" def __init__(self):\n",
" super().__init__()\n",
" self.theta = torch.nn.Parameter(torch.tensor(0.0))\n",
"\n",
" def forward(self, x):\n",
" return torch.cos(2 * 3.14 * x + self.theta)\n",
"\n",
"\n",
"mean = PriorMean().to(device)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "5",
"metadata": {},
"outputs": [],
"source": [
"with torch.no_grad():\n",
" res = y - mean(x).squeeze(-1)\n",
" m = mean(x_range)\n",
" res_quantiles = true_quantiles - m"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6",
"metadata": {},
"outputs": [
{
"data": {
"image/svg+xml": [
"\n",
"\n",
"\n"
],
"text/plain": [
""
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"fig, axes = plt.subplots(1, 2, figsize=(12, 4))\n",
"\n",
"axes[0].scatter(x.cpu(), y.cpu(), c=\"k\", marker=\".\")\n",
"axes[0].plot(x_range.cpu(), true_quantiles.cpu(), \"--\", c=\"gray\")\n",
"axes[0].plot(x_range.cpu(), m.detach().cpu())\n",
"\n",
"axes[1].scatter(x.cpu(), res.detach().cpu(), c=\"k\", marker=\".\")\n",
"axes[1].plot(x_range.cpu(), res_quantiles.detach().cpu(), \"--\", c=\"gray\")\n",
"\n",
"plt.show()"
]
},
{
"cell_type": "markdown",
"id": "7",
"metadata": {},
"source": [
"## Define models and likelihoods"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "8",
"metadata": {},
"outputs": [],
"source": [
"class MyGP(CenterGapQuantileGP):\n",
" def __init__(\n",
" self,\n",
" inducing_points,\n",
" num_quantiles,\n",
" num_lower_quantiles,\n",
" num_latents,\n",
" ):\n",
" N, D = inducing_points.size()\n",
" variational_distribution = CholeskyVariationalDistribution(\n",
" N,\n",
" batch_shape=torch.Size([num_latents]),\n",
" )\n",
" variational_strategy = LMCVariationalStrategy(\n",
" VariationalStrategy(\n",
" self,\n",
" inducing_points,\n",
" variational_distribution,\n",
" learn_inducing_locations=True,\n",
" ),\n",
" num_quantiles,\n",
" num_latents,\n",
" )\n",
"\n",
" mean = ConstantMean(batch_shape=torch.Size([num_latents]))\n",
" covar = ScaleKernel(\n",
" RBFKernel(ard_num_dims=D, batch_shape=torch.Size([num_latents])),\n",
" batch_shape=torch.Size([num_latents]),\n",
" )\n",
" super().__init__(\n",
" variational_strategy, mean, covar, [num_quantiles], [num_lower_quantiles]\n",
" )\n",
"\n",
"\n",
"inducing_points = torch.linspace(0, 1, 10).reshape(-1, 1).to(device)\n",
"central_q_index = (q - 0.5).abs().argmin().item()\n",
"num_latents = len(q) - 2\n",
"gp = MyGP(inducing_points, len(q), central_q_index, num_latents).to(device)\n",
"likelihood = CenterGapQuantileLikelihood(q, central_q_index).to(device)"
]
},
{
"cell_type": "markdown",
"id": "9",
"metadata": {},
"source": [
"## Train"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "10",
"metadata": {},
"outputs": [],
"source": [
"mean.train()\n",
"gp.train()\n",
"likelihood.train()\n",
"mll = VariationalELBO(likelihood, gp, num_data=y.numel())\n",
"optimizer = torch.optim.Adam(\n",
" list(mean.parameters()) + list(gp.parameters()) + list(likelihood.parameters()),\n",
" lr=0.001,\n",
")\n",
"\n",
"for _ in range(n_epochs):\n",
" optimizer.zero_grad()\n",
" res = y - mean(x).squeeze(-1)\n",
" output = gp(x)\n",
" loss = -mll(output, res)\n",
" loss.backward()\n",
" optimizer.step()"
]
},
{
"cell_type": "markdown",
"id": "11",
"metadata": {},
"source": [
"## Evaluate"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "12",
"metadata": {},
"outputs": [],
"source": [
"mean.eval()\n",
"gp.eval()\n",
"with torch.no_grad():\n",
" m = mean(x_pred)\n",
" mean_q = gp.mean_quantiles_mc(x_pred) + m\n",
" lower_q, upper_q = (\n",
" gp.quantile_quantiles_mc(x_pred, torch.tensor([0.025, 0.975]).to(device)) + m\n",
" )"
]
},
{
"cell_type": "markdown",
"id": "13",
"metadata": {},
"source": [
"## Plot result"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "14",
"metadata": {},
"outputs": [
{
"data": {
"image/svg+xml": [
"\n",
"\n",
"\n"
],
"text/plain": [
""
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"colors = plt.cm.tab10.colors\n",
"\n",
"plt.scatter(x.cpu(), y.cpu(), c=\"gray\", marker=\".\", alpha=0.1)\n",
"plt.plot(x_range.cpu(), true_quantiles.cpu(), \"--\", c=\"k\")\n",
"\n",
"for i in range(len(q)):\n",
" plt.plot(x_pred.cpu(), mean_q[:, i].cpu(), color=colors[i])\n",
" plt.fill_between(\n",
" x_pred.cpu().squeeze(),\n",
" lower_q[:, i].cpu(),\n",
" upper_q[:, i].cpu(),\n",
" color=colors[i],\n",
" alpha=0.3,\n",
" )\n",
"plt.show()"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "heavyedge",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.13.13"
}
},
"nbformat": 4,
"nbformat_minor": 5
}