Gaussian Process Regression#

[1]:
%pip install git+https://github.com/lightning-uq-box/lightning-uq-box.git
Collecting git+https://github.com/lightning-uq-box/lightning-uq-box.git
  Cloning https://github.com/lightning-uq-box/lightning-uq-box.git to /tmp/pip-req-build-9gywp8f9
  Running command git clone --filter=blob:none --quiet https://github.com/lightning-uq-box/lightning-uq-box.git /tmp/pip-req-build-9gywp8f9
  Resolved https://github.com/lightning-uq-box/lightning-uq-box.git to commit bb7f5757d8a82b4dd45347280551f01996c50da7
  Installing build dependencies ... - \ done
  Getting requirements to build wheel ... - done
  Preparing metadata (pyproject.toml) ... - done
Requirement already satisfied: einops>=0.3 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (0.8.2)
Requirement already satisfied: jsonargparse>=4.28.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from jsonargparse[signatures]>=4.28.0->lightning-uq-box==0.3.0.dev0) (4.46.0)
Requirement already satisfied: lightning>=2.4.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (2.6.1)
Requirement already satisfied: matplotlib>=3.5 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (3.10.8)
Requirement already satisfied: numpy>=1.21.1 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (1.26.4)
Requirement already satisfied: omegaconf>=2.3.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (2.3.0)
Requirement already satisfied: pandas>=1.1.3 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (3.0.0)
Requirement already satisfied: torch>=2.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (2.10.0)
Requirement already satisfied: torchmetrics>=1.2 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (1.8.2)
Requirement already satisfied: torchvision>=0.16.1 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (0.25.0)
Requirement already satisfied: scikit-learn>=1.3 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (1.8.0)
Requirement already satisfied: gpytorch>=1.11 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (1.15.1)
Requirement already satisfied: laplace-torch>=0.2.1 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (0.2.2.2)
Requirement already satisfied: curvlinops-for-pytorch==2.0.1 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (2.0.1)
Requirement already satisfied: uncertainty-toolbox>=0.1.1 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (0.1.1)
Requirement already satisfied: kornia>=0.6.9 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (0.8.2)
Requirement already satisfied: timm>=0.9.2 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (1.0.24)
Requirement already satisfied: torchseg>=0.0.1a1 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (0.0.1a4)
Requirement already satisfied: h5py>=3.12.1 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (3.15.1)
Requirement already satisfied: ema-pytorch>=0.7.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (0.7.9)
Requirement already satisfied: vbll>=0.4.9 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-uq-box==0.3.0.dev0) (0.4.9)
Requirement already satisfied: backpack-for-pytorch<2.0.0,>=1.6.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from curvlinops-for-pytorch==2.0.1->lightning-uq-box==0.3.0.dev0) (1.7.1)
Requirement already satisfied: scipy<2.0.0,>=1.7.1 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from curvlinops-for-pytorch==2.0.1->lightning-uq-box==0.3.0.dev0) (1.17.0)
Requirement already satisfied: tqdm<5.0.0,>=4.61.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from curvlinops-for-pytorch==2.0.1->lightning-uq-box==0.3.0.dev0) (4.67.3)
Requirement already satisfied: einconv in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from curvlinops-for-pytorch==2.0.1->lightning-uq-box==0.3.0.dev0) (0.1.0)
Requirement already satisfied: unfoldNd<1.0.0,>=0.2.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from backpack-for-pytorch<2.0.0,>=1.6.0->curvlinops-for-pytorch==2.0.1->lightning-uq-box==0.3.0.dev0) (0.2.3)
Requirement already satisfied: packaging in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from unfoldNd<1.0.0,>=0.2.0->backpack-for-pytorch<2.0.0,>=1.6.0->curvlinops-for-pytorch==2.0.1->lightning-uq-box==0.3.0.dev0) (26.0)
Requirement already satisfied: jaxtyping in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from gpytorch>=1.11->lightning-uq-box==0.3.0.dev0) (0.3.7)
Requirement already satisfied: mpmath<=1.3,>=0.19 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from gpytorch>=1.11->lightning-uq-box==0.3.0.dev0) (1.3.0)
Requirement already satisfied: linear_operator>=0.6 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from gpytorch>=1.11->lightning-uq-box==0.3.0.dev0) (0.6)
Requirement already satisfied: PyYAML>=3.13 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from jsonargparse>=4.28.0->jsonargparse[signatures]>=4.28.0->lightning-uq-box==0.3.0.dev0) (6.0.3)
Requirement already satisfied: docstring-parser>=0.17 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from jsonargparse[signatures]>=4.28.0->lightning-uq-box==0.3.0.dev0) (0.17.0)
Requirement already satisfied: typeshed-client>=2.8.2 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from jsonargparse[signatures]>=4.28.0->lightning-uq-box==0.3.0.dev0) (2.8.2)
Requirement already satisfied: kornia_rs>=0.1.9 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from kornia>=0.6.9->lightning-uq-box==0.3.0.dev0) (0.1.10)
Requirement already satisfied: asdfghjkl==0.1a4 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from laplace-torch>=0.2.1->lightning-uq-box==0.3.0.dev0) (0.1a4)
Requirement already satisfied: opt_einsum in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from laplace-torch>=0.2.1->lightning-uq-box==0.3.0.dev0) (3.4.0)
Requirement already satisfied: fsspec<2028.0,>=2022.5.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from fsspec[http]<2028.0,>=2022.5.0->lightning>=2.4.0->lightning-uq-box==0.3.0.dev0) (2026.1.0)
Requirement already satisfied: lightning-utilities<2.0,>=0.10.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning>=2.4.0->lightning-uq-box==0.3.0.dev0) (0.15.2)
Requirement already satisfied: typing-extensions<6.0,>4.5.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning>=2.4.0->lightning-uq-box==0.3.0.dev0) (4.15.0)
Requirement already satisfied: pytorch-lightning in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning>=2.4.0->lightning-uq-box==0.3.0.dev0) (2.6.1)
Requirement already satisfied: aiohttp!=4.0.0a0,!=4.0.0a1 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from fsspec[http]<2028.0,>=2022.5.0->lightning>=2.4.0->lightning-uq-box==0.3.0.dev0) (3.13.3)
Requirement already satisfied: setuptools in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from lightning-utilities<2.0,>=0.10.0->lightning>=2.4.0->lightning-uq-box==0.3.0.dev0) (80.9.0)
Requirement already satisfied: filelock in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (3.20.3)
Requirement already satisfied: sympy>=1.13.3 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (1.14.0)
Requirement already satisfied: networkx>=2.5.1 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (3.6.1)
Requirement already satisfied: jinja2 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (3.1.6)
Requirement already satisfied: cuda-bindings==12.9.4 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (12.9.4)
Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.8.93 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (12.8.93)
Requirement already satisfied: nvidia-cuda-runtime-cu12==12.8.90 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (12.8.90)
Requirement already satisfied: nvidia-cuda-cupti-cu12==12.8.90 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (12.8.90)
Requirement already satisfied: nvidia-cudnn-cu12==9.10.2.21 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (9.10.2.21)
Requirement already satisfied: nvidia-cublas-cu12==12.8.4.1 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (12.8.4.1)
Requirement already satisfied: nvidia-cufft-cu12==11.3.3.83 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (11.3.3.83)
Requirement already satisfied: nvidia-curand-cu12==10.3.9.90 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (10.3.9.90)
Requirement already satisfied: nvidia-cusolver-cu12==11.7.3.90 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (11.7.3.90)
Requirement already satisfied: nvidia-cusparse-cu12==12.5.8.93 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (12.5.8.93)
Requirement already satisfied: nvidia-cusparselt-cu12==0.7.1 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (0.7.1)
Requirement already satisfied: nvidia-nccl-cu12==2.27.5 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (2.27.5)
Requirement already satisfied: nvidia-nvshmem-cu12==3.4.5 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (3.4.5)
Requirement already satisfied: nvidia-nvtx-cu12==12.8.90 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (12.8.90)
Requirement already satisfied: nvidia-nvjitlink-cu12==12.8.93 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (12.8.93)
Requirement already satisfied: nvidia-cufile-cu12==1.13.1.3 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (1.13.1.3)
Requirement already satisfied: triton==3.6.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from torch>=2.0->lightning-uq-box==0.3.0.dev0) (3.6.0)
Requirement already satisfied: cuda-pathfinder~=1.1 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from cuda-bindings==12.9.4->torch>=2.0->lightning-uq-box==0.3.0.dev0) (1.3.3)
Requirement already satisfied: aiohappyeyeballs>=2.5.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2028.0,>=2022.5.0->lightning>=2.4.0->lightning-uq-box==0.3.0.dev0) (2.6.1)
Requirement already satisfied: aiosignal>=1.4.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2028.0,>=2022.5.0->lightning>=2.4.0->lightning-uq-box==0.3.0.dev0) (1.4.0)
Requirement already satisfied: attrs>=17.3.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2028.0,>=2022.5.0->lightning>=2.4.0->lightning-uq-box==0.3.0.dev0) (25.4.0)
Requirement already satisfied: frozenlist>=1.1.1 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2028.0,>=2022.5.0->lightning>=2.4.0->lightning-uq-box==0.3.0.dev0) (1.8.0)
Requirement already satisfied: multidict<7.0,>=4.5 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2028.0,>=2022.5.0->lightning>=2.4.0->lightning-uq-box==0.3.0.dev0) (6.7.1)
Requirement already satisfied: propcache>=0.2.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2028.0,>=2022.5.0->lightning>=2.4.0->lightning-uq-box==0.3.0.dev0) (0.4.1)
Requirement already satisfied: yarl<2.0,>=1.17.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2028.0,>=2022.5.0->lightning>=2.4.0->lightning-uq-box==0.3.0.dev0) (1.22.0)
Requirement already satisfied: idna>=2.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from yarl<2.0,>=1.17.0->aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]<2028.0,>=2022.5.0->lightning>=2.4.0->lightning-uq-box==0.3.0.dev0) (3.11)
Requirement already satisfied: contourpy>=1.0.1 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from matplotlib>=3.5->lightning-uq-box==0.3.0.dev0) (1.3.3)
Requirement already satisfied: cycler>=0.10 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from matplotlib>=3.5->lightning-uq-box==0.3.0.dev0) (0.12.1)
Requirement already satisfied: fonttools>=4.22.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from matplotlib>=3.5->lightning-uq-box==0.3.0.dev0) (4.61.1)
Requirement already satisfied: kiwisolver>=1.3.1 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from matplotlib>=3.5->lightning-uq-box==0.3.0.dev0) (1.4.9)
Requirement already satisfied: pillow>=8 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from matplotlib>=3.5->lightning-uq-box==0.3.0.dev0) (12.1.0)
Requirement already satisfied: pyparsing>=3 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from matplotlib>=3.5->lightning-uq-box==0.3.0.dev0) (3.3.2)
Requirement already satisfied: python-dateutil>=2.7 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from matplotlib>=3.5->lightning-uq-box==0.3.0.dev0) (2.9.0.post0)
Requirement already satisfied: antlr4-python3-runtime==4.9.* in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from omegaconf>=2.3.0->lightning-uq-box==0.3.0.dev0) (4.9.3)
Requirement already satisfied: six>=1.5 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from python-dateutil>=2.7->matplotlib>=3.5->lightning-uq-box==0.3.0.dev0) (1.17.0)
Requirement already satisfied: joblib>=1.3.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from scikit-learn>=1.3->lightning-uq-box==0.3.0.dev0) (1.5.3)
Requirement already satisfied: threadpoolctl>=3.2.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from scikit-learn>=1.3->lightning-uq-box==0.3.0.dev0) (3.6.0)
Requirement already satisfied: huggingface_hub in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from timm>=0.9.2->lightning-uq-box==0.3.0.dev0) (1.3.7)
Requirement already satisfied: safetensors in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from timm>=0.9.2->lightning-uq-box==0.3.0.dev0) (0.7.0)
Requirement already satisfied: importlib_resources>=1.4.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from typeshed-client>=2.8.2->jsonargparse[signatures]>=4.28.0->lightning-uq-box==0.3.0.dev0) (6.5.2)
Requirement already satisfied: hf-xet<2.0.0,>=1.2.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from huggingface_hub->timm>=0.9.2->lightning-uq-box==0.3.0.dev0) (1.2.0)
Requirement already satisfied: httpx<1,>=0.23.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from huggingface_hub->timm>=0.9.2->lightning-uq-box==0.3.0.dev0) (0.28.1)
Requirement already satisfied: shellingham in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from huggingface_hub->timm>=0.9.2->lightning-uq-box==0.3.0.dev0) (1.5.4)
Requirement already satisfied: typer-slim in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from huggingface_hub->timm>=0.9.2->lightning-uq-box==0.3.0.dev0) (0.21.1)
Requirement already satisfied: anyio in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from httpx<1,>=0.23.0->huggingface_hub->timm>=0.9.2->lightning-uq-box==0.3.0.dev0) (4.12.1)
Requirement already satisfied: certifi in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from httpx<1,>=0.23.0->huggingface_hub->timm>=0.9.2->lightning-uq-box==0.3.0.dev0) (2026.1.4)
Requirement already satisfied: httpcore==1.* in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from httpx<1,>=0.23.0->huggingface_hub->timm>=0.9.2->lightning-uq-box==0.3.0.dev0) (1.0.9)
Requirement already satisfied: h11>=0.16 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from httpcore==1.*->httpx<1,>=0.23.0->huggingface_hub->timm>=0.9.2->lightning-uq-box==0.3.0.dev0) (0.16.0)
Requirement already satisfied: wadler-lindig>=0.1.3 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from jaxtyping->gpytorch>=1.11->lightning-uq-box==0.3.0.dev0) (0.1.7)
Requirement already satisfied: MarkupSafe>=2.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from jinja2->torch>=2.0->lightning-uq-box==0.3.0.dev0) (3.0.3)
Requirement already satisfied: click>=8.0.0 in /home/docs/checkouts/readthedocs.org/user_builds/lightning-uq-box/envs/latest/lib/python3.12/site-packages (from typer-slim->huggingface_hub->timm>=0.9.2->lightning-uq-box==0.3.0.dev0) (8.3.1)
Note: you may need to restart the kernel to use updated packages.

Theoretic Foundation#

The goal of previously introduced methods was to find a distribution over the weights of a parameterized function i.e. a neural network. In contrast, the basic idea of a Gaussian Process (GP) is to instead consider a distribution over possible functions, that fit the data in some way. Formally,

“A Gaussian process is a collection of random variables, any finite number of which have a joint Gaussian distribution.” Seeger, 2004

Precisely, a GP can be described by a possibly infinite amount of function values

\[f(x) \sim \mathcal{GP}(m(x),k_{\gamma}(x)),\]

such that any finite collection of function values \(f\) has a joint Gaussian distribution,

\[f = f(X) = [f(x_1),\dots,f(x_K)]^{\top} \sim \mathcal{N}(m_X,\mathcal{K}_{X,X}),\]

with a mean vector, \((m_X)i = m(x_i)\), and covariance matrix, \((\mathcal{K}_{X,X})_{ij} = k_{\gamma}(x_i,x_j)\), stemming from the mean function \(m\) and covariance kernel of the GP, \(k_{\gamma}\), that is parametrized by \(\gamma\). A commonly used covariance function is the squared exponential, also referred to as Radial Basis Function (RBF) kernel, exponentiated quadratic or Gaussian kernel:

\[k_{\gamma}(x,x') = \text{cov}(f(x),f(x')) = \eta^2\exp{\left(-\frac{1}{2l^2}|x-x'|^2\right)}.\]

Where \(\gamma =( \eta^2, l)\) and \(\eta^2\) can be set to \(1\) or tuned as a hyperparameter. By default the lengthscale \(l=1\) but can also be optimized over. Now the GP, \(f (x) \sim GP(m(x), k(x, x'))\), as a distribution over functions can be used to solve a regression problem. Following Seeger, 2004, consider the simple special case where the observations are noise free and you have training data \(\mathcal{D}_{\mathrm{train}} = \{(x_i, y_i)\}_{i=1}^{K}\) with \(X=(x_i)_{i=1}^K\) and \(Y=(y_i)_{i=1}^K\). The joint prior distribution of the training outputs, \(Y\), and the test outputs \(f_*=f_*(X_*)= (f(x_k))_{k=1}^m\) where \(X_* = (x_k)_{k=1}^m\) are the test points, according to the prior is

\[\begin{split} p(Y,f_*) = \mathcal{N}\bigg(0, \Bigg[ \begin{array}{rr} \mathcal{K}_{X,X} & \mathcal{K}_{X,X_*} \\ \mathcal{K}_{X_*,X} & \mathcal{K}_{X_*,X_*} \\ \end{array} \Bigg] \bigg).\end{split}\]

Here the mean function is assumed to be \(m_X = 0\) and \(\mathcal{K}_{X,X_*}\) denotes the \(K \times m\) matrix of the covariances evaluated at all pairs of training and test points, and similarly for the other entries \(\mathcal{K}_{X,X}\), \(\mathcal{K}_{X_*,X_*}\) and \(\mathcal{K}_{X_*,X}\). To make predictions based on the knowledge of the training points, conditioning on the prior observations is used and yields,

\[\begin{split}p(f_*|X_*, X, Y) = \mathcal{N}(\mathcal{K}_{X_*,X}\mathcal{K}_{X,X}^{-1}Y, \mathcal{K}_{X_*,X_*}-\mathcal{K}_{X_*,X}\mathcal{K}_{X,X}^{-1}\mathcal{K}_{X_*,X})\\ = \mathcal{N}(m(X,X_*,Y), \tilde{\mathcal{K}}_{X,X_*}).\end{split}\]

Now to generate function values on test points, one uses samples from the posterior distribution \(f_*(X_*) \sim \mathcal{N}(m(X,X_*,Y), \tilde{K}(X, X_*))\). To illustrate how we can obtain these samples from the posterior distribution, consider a Gaussian with arbitrary mean \(m\) and covariance \(K\), i.e. \(f_* \sim \mathcal{N}(m,K)\). For this one can use a scalar Gaussian generator, which is available in many packages:

  • Compute the Cholesky decomposition of \(K=LL^T\), where \(L\) is a lower triangular matrix. This works because \(K\) is symmetric by definition.

  • Then, draw multiple \(u \sim \mathcal{N}(0, I)\).

  • Now, compute the samples with \(f_* = m + Lu\). This has the desired mean, \(m\) and covariance \(L\mathbb{E}(uu^T)L^T = LL^T = K\).

The above can be extended to incorporate noisy measurements \(y \rightarrow y + e\), see Seeger, 2004, or noise on the inputs as in Johnson, 2019. Both of these extensions require tuning of further hyperparameters, yet beneficially allow to incorporate a prediction of aleatoric uncertainty in a GP.

For example, assume additive Gaussian noise on the distribution of the function values,

\[p(y(x)|f(x)) = \mathcal{N}(y(x); f(x),\sigma^2).\]

Then the predictive distribution of the GP evaluated at the \(K_*\) test points, \(X_*\), is given by

\[p(f_*|X_*,X,Y,\gamma,\sigma^2) = \mathcal{N}(\mathbb{E}[f_*],\text{cov}(f_*)),\]
\[\mathbb{E}[f_*] = m_{X_*} + \mathcal{K}_{X_*,X}[\mathcal{K}_{X,X}+\sigma^2 I]^{-1}Y,\]
\[\text{cov}(f_*) = \mathcal{K}_{X_*,X_*} - \mathcal{K}_{X_*,X}[\mathcal{K}_{X,X}+\sigma^2 I]^{-1}\mathcal{K}_{X,X_*}.\]

Here \(m_{X_*}\) is the \(K_* \times 1\) mean vector, which is assumed to be zero in the previous case.

In both cases, with and without additive noise on the function values, the GP is trained by learning interpretable kernel hyperparameters. The log marginal likelihood of the targets \(y\) - the probability of the data conditioned only on kernel hyperparameters \(\gamma\) - provides a principled probabilistic framework for kernel learning:

\[\log p(y | \gamma, X) \propto -\left(y^{\top}(\mathcal{K}_{\gamma}+\sigma^2 I)^{-1}y + \log|\mathcal{K}_{\gamma} + \sigma^2 I|\right),\]

where \(\mathcal{K}_{\gamma}\) is used for \(\mathcal{K}_{X,X}\) given \(\gamma\). Kernel learning can be achieved by optimizing the above with respect to \(\gamma\).

The computational bottleneck for inference is solving the linear system \((\mathcal{K}_{X,X}+\sigma^2 I)^{-1}y\), and for kernel learning it is computing the log determinant \(\log|\mathcal{K}_{X,X}+ \sigma^2 I|\) in the marginal likelihood. The standard approach is to compute the Cholesky decomposition of the \(K \times K\) matrix \(\mathcal{K}_{X,X}\), which requires \(\mathcal{O}(K^3)\) operations and \(\mathcal{O}(K^2)\) storage. After inference is complete, the predictive mean costs \(\mathcal{O}(K)\), and the predictive variance costs \(\mathcal{O}(K^2)\), per test point \(x_*\).

Imports#

[2]:
import tempfile

import gpytorch
import matplotlib.pyplot as plt
import numpy as np
import torch
from lightning.pytorch import seed_everything
from tqdm.notebook import tqdm

from lightning_uq_box.datamodules import ToyHeteroscedasticDatamodule
from lightning_uq_box.viz_utils import (
    plot_calibration_uq_toolbox,
    plot_predictions_regression,
    plot_toy_regression_data,
)

plt.rcParams["figure.figsize"] = [14, 5]
[3]:
seed_everything(0)  # seed everything for reproducibility
Seed set to 0
[3]:
0

We define a temporary directory to look at some training metrics and results.

[4]:
my_temp_dir = tempfile.mkdtemp()

Datamodule#

To demonstrate the method, we will make use of a Toy Regression Example that is defined as a Lightning Datamodule. While this might seem like overkill for a small toy problem, we think it is more helpful how the individual pieces of the library fit together so you can train models on more complex tasks.

[5]:
dm = ToyHeteroscedasticDatamodule()

X_train, Y_train, train_loader, X_test, Y_test, test_loader, X_gtext, Y_gtext = (
    dm.X_train,
    dm.Y_train,
    dm.train_dataloader(),
    dm.X_test,
    dm.Y_test,
    dm.test_dataloader(),
    dm.X_gtext,
    dm.Y_gtext,
)
[6]:
fig = plot_toy_regression_data(X_train, Y_train, X_test, Y_test)
../../_images/tutorials_regression_gp_10_0.png

Model#

[7]:
class ExactGPModel(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        super().__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ConstantMean()
        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel())

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)
[8]:
# initialize the likelihood
likelihood = gpytorch.likelihoods.GaussianLikelihood()

# init the GP model
gp_model = ExactGPModel(X_train.squeeze(), Y_train.squeeze(), likelihood)

# Find optimal model hyperparameters
gp_model.train()
likelihood.train()

# Use the adam optimizer
optimizer = torch.optim.Adam(
    gp_model.parameters(), lr=1e-2
)  # Includes GaussianLikelihood parameters

# "Loss" for GPs - the marginal log likelihood
mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, gp_model)

bar = tqdm(range(100))
for i in bar:
    # Zero gradients from previous iteration
    optimizer.zero_grad()
    # Output from model
    output = gp_model(X_train.squeeze())
    # Calc loss and backprop gradients
    loss = -mll(output, Y_train.squeeze())
    loss.backward()
    optimizer.step()
    bar.set_postfix(loss=f"{loss.detach().cpu().item()}")

Prediction#

[9]:
gp_model.eval()
likelihood.eval()

with torch.no_grad():
    gp_preds = gp_model(X_gtext.cpu())

gp_mean = gp_preds.mean.detach().cpu().numpy()
gp_var = gp_preds.variance.detach().cpu().numpy()
gp_covar = gp_preds.covariance_matrix.detach().cpu().numpy()

Evaluate Predictions#

[10]:
fig = plot_predictions_regression(
    X_train,
    Y_train,
    X_gtext,
    Y_gtext,
    gp_mean[:, None],
    np.sqrt(gp_var),
    epistemic=np.sqrt(gp_var),
    title="Gaussian Process",
    show_bands=False,
)
../../_images/tutorials_regression_gp_17_0.png

For some additional metrics relevant to UQ, we can use the great uncertainty-toolbox that gives us some insight into the calibration of our prediction. For a discussion of why this is important, see …

[11]:
with torch.no_grad():
    gp_preds = gp_model(X_test.cpu())

gp_mean = gp_preds.mean.detach().cpu().numpy()
gp_var = gp_preds.variance.detach().cpu().numpy()
gp_covar = gp_preds.covariance_matrix.detach().cpu().numpy()

fig = plot_calibration_uq_toolbox(
    gp_mean, np.sqrt(gp_var), Y_test.cpu().numpy(), X_test.cpu().numpy()
)
../../_images/tutorials_regression_gp_19_0.png