Overview

GeometricKernels is a library that implements natural kernels (heat [1], Matérn) on such non-Euclidean spaces as Riemannian manifolds, graphs and meshes.

The main projected application is Gaussian processes.

Installation and Requirements

This is a Python 3 library.

Before doing anything, you might want to create and activate a new virtual environment:

uv

uv venv --python python[version] [venv_dir]

where [env_dir] is the directory (default is .venv) of the environment and [version] is the version of Python you want to use, we currently support 3.9, 3.10, 3.11, 3.12.

[Optional] activate the environment. However, this is not strictly necessary for uv. Instead, use tools like uv run python to run Python inside the environment. See uv documentation <https://docs.astral.sh/uv/> for more details.

Conda

conda create -n [env_name] python=[version]
conda activate [env_name]

where [env_name] is the name of the environment and [version] is the version of Python you want to use, we currently support 3.9, 3.10, 3.11, 3.12.

Virtualenv

virtualenv [env_name]
source [env_name]/bin/activate

To install GeometricKernels, run

pip install geometric_kernels

Note

If you use uv, swap pip with uv pip everywhere. Additionally, if you ran uv init, use uv add to add the requirement in your pyproject.toml and install it.

Note

If you want to install specific GitHub branch called [branch], run

pip install "git+https://github.com/geometric-kernels/GeometricKernels@[branch]"

The kernels are compatible with several backends, namely

• NumPy,

• TensorFlow (can be used together with the GP library GPflow),

• PyTorch (can be used together with the GP library GPyTorch),

• JAX (can be used together with the GP library GPJax).

Any backend, except for NumPy, should be manually installed.

TensorFlow Installation

You need both tensorflow and tensorflow-probability. You can get them by running

pip install tensorflow tensorflow-probability

[Optional] We support the TensorFlow-based Gaussian process library gpflow which you can install by running

pip install gpflow

PyTorch Installation

You can get PyTorch by running

pip install torch

[Optional] We support the PyTorch-based Gaussian process library gpytorch which you can install by running

pip install gpytorch

JAX Installation

To install JAX, follow these instructions.

[Optional] We support the JAX-based Gaussian process library GPJax which you can install by running

pip install gpjax

A Basic Example

In the following example we show how to initialize the Matern52 kernel on the two-dimensional sphere and how to compute a kernel matrix for a few points on the sphere.

Numpy

>>> # Import a backend.
>>> import numpy as np
>>> # Import the geometric_kernels backend.
>>> import geometric_kernels
>>> # Import a space and an appropriate kernel.
>>> from geometric_kernels.spaces import Hypersphere
>>> from geometric_kernels.kernels import MaternGeometricKernel

>>> # Create a manifold (2-dim sphere).
>>> hypersphere = Hypersphere(dim=2)

>>> # Define 3 points on the sphere.
>>> xs = np.array([[0., 0., 1.], [0., 1., 0.], [1., 0., 0.]])

>>> # Initialize kernel.
>>> kernel = MaternGeometricKernel(hypersphere)
>>> params = kernel.init_params()
>>> params["nu"] = np.array([5/2])
>>> params["lengthscale"] = np.array([1.])

>>> # Compute and print out the 3x3 kernel matrix.
>>> print(np.around(kernel.K(params, xs), 2))
[[1.   0.36 0.36]
 [0.36 1.   0.36]
 [0.36 0.36 1.  ]]

TensorFlow

>>> import numpy as np
>>> # Import a backend.
>>> import tensorflow as tf
>>> # Import the geometric_kernels backend.
>>> import geometric_kernels.tensorflow
>>> # Import a space and an appropriate kernel.
>>> from geometric_kernels.spaces import Hypersphere
>>> from geometric_kernels.kernels import MaternGeometricKernel

>>> # Create a manifold (2-dim sphere).
>>> hypersphere = Hypersphere(dim=2)

>>> # Define 3 points on the sphere.
>>> xs = tf.convert_to_tensor([[0., 0., 1.], [0., 1., 0.], [1., 0., 0.]])

>>> # Initialize kernel.
>>> kernel = MaternGeometricKernel(hypersphere)
>>> params = kernel.init_params()
>>> params["nu"] = tf.convert_to_tensor([5/2])
>>> params["lengthscale"] = tf.convert_to_tensor([1.])

>>> # Compute and print out the 3x3 kernel matrix.
>>> print(np.around(kernel.K(params, xs).numpy(), 2))
[[1.   0.36 0.36]
 [0.36 1.   0.36]
 [0.36 0.36 1.  ]]

PyTorch

>>> import numpy as np
>>> # Import a backend.
>>> import torch
>>> # Import the geometric_kernels backend.
>>> import geometric_kernels.torch
>>> # Import a space and an appropriate kernel.
>>> from geometric_kernels.spaces import Hypersphere
>>> from geometric_kernels.kernels import MaternGeometricKernel

>>> # Create a manifold (2-dim sphere).
>>> hypersphere = Hypersphere(dim=2)

>>> # Define 3 points on the sphere.
>>> xs = torch.tensor([[0., 0., 1.], [0., 1., 0.], [1., 0., 0.]])

>>> # Initialize kernel.
>>> kernel = MaternGeometricKernel(hypersphere)
>>> params = kernel.init_params()
>>> params["nu"] = torch.tensor([5/2])
>>> params["lengthscale"] = torch.tensor([1.])

>>> # Compute and print out the 3x3 kernel matrix.
>>> print(np.around(kernel.K(params, xs).detach().cpu().numpy(), 2))
[[1.   0.36 0.36]
 [0.36 1.   0.36]
 [0.36 0.36 1.  ]]

JAX

>>> import numpy as np
>>> # Import a backend.
>>> import jax.numpy as jnp
>>> # Import the geometric_kernels backend.
>>> import geometric_kernels.jax
>>> # Import a space and an appropriate kernel.
>>> from geometric_kernels.spaces import Hypersphere
>>> from geometric_kernels.kernels import MaternGeometricKernel

>>> # Create a manifold (2-dim sphere).
>>> hypersphere = Hypersphere(dim=2)

>>> # Define 3 points on the sphere.
>>> xs = jnp.array([[0., 0., 1.], [0., 1., 0.], [1., 0., 0.]])

>>> # Initialize kernel.
>>> kernel = MaternGeometricKernel(hypersphere)
>>> params = kernel.init_params()
>>> params["nu"] = jnp.array([5/2])
>>> params["lengthscale"] = jnp.array([1.0])

>>> # Compute and print out the 3x3 kernel matrix.
>>> print(np.around(np.asarray(kernel.K(params, xs)), 2))
[[1.   0.36 0.36]
 [0.36 1.   0.36]
 [0.36 0.36 1.  ]]

You can find more examples here.

Citation

If you are using GeometricKernels, please cite the library paper:

@article{mostowsky2024,
      title = {The GeometricKernels Package: Heat and Matérn Kernels for Geometric Learning on Manifolds, Meshes, and Graphs},
      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},
      year = {2024},
      journal = {arXiv:2407.08086},
}

Please also consider citing the theoretical papers the library is based on. You can find the relevant references for any space in

• the docstring of the respective space class,

• at the end of the respective tutorial notebook.

Footnotes

[1]

The heat kernel (or diffusion kernel) is a far-reaching generalization of the RBF kernel (a.k.a. Gaussian kernel, or squared exponential kernel). It can be considered to be a Matérn kernel with smoothness parameter \(\nu = \infty\), as we do in this library.