geometric-kernels/GeometricKernels

GeometricKernels – Kernels on non‑Euclidean spaces

What it is – A Python library that implements classic kernel functions (heat, Matérn, etc.) for data that lives on geometric domains such as Riemannian manifolds, graphs, and triangle meshes. By providing these kernels, the package lets you plug Gaussian‑process (GP) models into spaces where ordinary Euclidean kernels don’t apply.

Why it matters – Many modern ML tasks involve structured domains (e.g., sensor networks on road graphs, shape analysis on meshes, or data on spheres). Standard GP toolkits assume flat Euclidean inputs, so they cannot capture the intrinsic geometry. GeometricKernels supplies the mathematically‑correct covariance functions, enabling uncertainty‑aware learning on those domains.

Quick start (installation & a tiny example)

# optional: create a virtual environment (uv, conda, or virtualenv)
pip install geometric_kernels          # core package
# pick one backend – only the one you need
pip install torch                     # PyTorch backend (or tensorflow, jax, etc.)

import numpy as np
import geometric_kernels as gk
from geometric_kernels.spaces import Hypersphere
from geometric_kernels.kernels import MaternGeometricKernel

# 2‑sphere (dim=2)
S2 = Hypersphere(dim=2)

# three points on the sphere (Cartesian coordinates)
X = np.array([[0., 0., 1.],
              [0., 1., 0.],
              [1., 0., 0.]])

k = MaternGeometricKernel(S2)          # Matérn‑5/2 kernel on the sphere
params = k.init_params()
params["nu"] = np.array([5/2])
params["lengthscale"] = np.array([1.])

print(np.around(k.K(params, X), 2))
# → [[1.   0.36 0.36]
#    [0.36 1.   0.36]
#    [0.36 0.36 1.  ]]

The snippet shows how a 3×3 covariance matrix is built for points on a sphere using the NumPy backend.

How it fits into the ML ecosystem

Layer	What GeometricKernels provides
Domain	Abstract space classes (e.g., `Hypersphere`, `Graph`, `Mesh`). They know distances, geodesics, and volume elements on the underlying geometry.
Kernels	Heat kernel, Matérn kernels (ν = 1/2, 3/2, 5/2, …) that respect the geometry of the space.
Backends	Thin wrapper (`lab`) that can run on NumPy, TensorFlow, PyTorch, or JAX. Choose the one that matches your downstream code.
GP front‑ends	Small adapters for popular GP libraries – GPflow (TensorFlow), GPyTorch (PyTorch), GPJax (JAX). They expose the kernels as drop‑in replacements for the libraries’ own covariance functions.
Applications	Graph‑structured regression with uncertainty, Bayesian optimisation on manifolds, shape modelling, geophysical data on the Earth sphere, etc.

Documentation & learning resources

Website – https://geometric‑kernels.github.io/GeometricKernels – full API reference, tutorials, and a searchable index.
Example notebooks – hosted under the repo’s notebooks/ folder and linked from the docs. They cover every supported space, backend, and GP front‑end, plus a custom‑space tutorial.
Benchmarks – the PeMS Regression benchmark demonstrates graph‑node regression with uncertainty and shows GeometricKernels‑based GPs beating GNN ensembles.
Bayesian optimisation demo – a minimal notebook using botorch to illustrate BO on a manifold.

Development workflow (if you want to contribute)

Clone the repo and run make venv (or set up your own env).
make install pulls all optional backends and dev dependencies.
Run make lint and make test to check style and unit tests.
Follow the contribution guide – a PR example is linked in the README.

Citation

If you use the library in research, cite the JMLR paper:

@article{JMLR:v26:24-1185,
  author  = {Peter Mostowsky and Vincent Dutordoir and Iskander Azangulov and Noémie Jaquier and Michael John Hutchinson and Aditya Ravuri and Leonel Rozo and Alexander Terenin and Viacheslav Borovitskiy},
  title   = {The GeometricKernels Package: Heat and Matérn Kernels for Geometric Learning on Manifolds, Meshes, and Graphs},
  journal = {Journal of Machine Learning Research},
  year    = {2025},
  volume  = {26},
  number  = {276},
  pages   = {1--14},
  url     = {http://jmlr.org/papers/v26/24-1185.html}
}

Additional references are listed in each space’s docstring and in the tutorial notebooks.

TL;DR

GeometricKernels is a well‑maintained, backend‑agnostic Python package that brings mathematically sound kernels to manifolds, graphs, and meshes, enabling Gaussian‑process models (and downstream Bayesian methods) to work on non‑Euclidean data. Install it, pick a backend, and you can immediately use the kernels inside GPflow, GPyTorch, or GPJax.