diff --git a/README.md b/README.md index f3b18e3e..1bf16b4e 100644 --- a/README.md +++ b/README.md @@ -161,6 +161,8 @@ Run the tests make test ``` +**Example:** If you want to learn how to implement your own space or kernel component, checkout the [CustomSpacesAndKernels.ipynb](https://github.com/geometric-kernels/GeometricKernels/blob/main/notebooks/CustomSpacesAndKernels.ipynb) notebook. + ## If you have a question Post it in issues using the `"How do I do ..." and other issues` template and the "question" label. diff --git a/docs/examples/CustomSpacesAndKernels.nblink b/docs/examples/CustomSpacesAndKernels.nblink new file mode 100644 index 00000000..d8e53e42 --- /dev/null +++ b/docs/examples/CustomSpacesAndKernels.nblink @@ -0,0 +1 @@ +{"path": "../../notebooks/CustomSpacesAndKernels.ipynb"} \ No newline at end of file diff --git a/docs/examples/examples.rst b/docs/examples/examples.rst index df7a60b5..e65f501a 100644 --- a/docs/examples/examples.rst +++ b/docs/examples/examples.rst @@ -6,6 +6,7 @@ Spaces .. toctree:: :titlesonly: + Custom Spaces and Kernels Graph GraphEdges Hyperbolic diff --git a/docs/examples/introduction.rst b/docs/examples/introduction.rst index 5e48d827..9e6a8fb2 100644 --- a/docs/examples/introduction.rst +++ b/docs/examples/introduction.rst @@ -6,6 +6,8 @@ If you are interested in using GeometricKernels with a backend other than NumPy, If you want to use GeometricKernels with one of the popular Gaussian process libraries, checkout the *Frontends* notebooks. +If you want to learn how to implement your own space or kernel component, checkout `CustomSpacesAndKernels.ipynb `_ notebook. + Finally, if you are interested in application examples: * `PeMS Regression `__ diff --git a/docs/index.rst b/docs/index.rst index b0bcb630..4069ae4b 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -395,9 +395,18 @@ You can find more examples :doc:`here `. Citation ======== -If you are using GeometricKernels, please consider citing the theoretical papers it is based on. +If you are using GeometricKernels, please cite the `library paper `__: -You can find the relevant references for any space in +.. code-block:: latex + + @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. diff --git a/geometric_kernels/spaces/graph_edges.py b/geometric_kernels/spaces/graph_edges.py index 5df4c39d..b712b47b 100644 --- a/geometric_kernels/spaces/graph_edges.py +++ b/geometric_kernels/spaces/graph_edges.py @@ -4,9 +4,8 @@ import lab as B import numpy as np -from beartype.door import is_bearable from beartype.typing import Dict, List, Optional, Tuple, Union -from scipy.sparse import csr_matrix, lil_matrix +from scipy.sparse import csr_matrix, issparse, lil_matrix from geometric_kernels.lab_extras import ( SparseArray, @@ -447,7 +446,7 @@ def from_adjacency( # noqa: C901 """ if isinstance(adjacency_matrix, np.ndarray): index = csr_matrix(adjacency_matrix, dtype=int) - elif is_bearable(adjacency_matrix, SparseArray): + elif issparse(adjacency_matrix): index = csr_matrix(adjacency_matrix, dtype=int, copy=True) else: raise ValueError( diff --git a/notebooks/CustomSpacesAndKernels.ipynb b/notebooks/CustomSpacesAndKernels.ipynb new file mode 100644 index 00000000..2878836c --- /dev/null +++ b/notebooks/CustomSpacesAndKernels.ipynb @@ -0,0 +1,756 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "cc7f6575-7e7e-419e-80ae-81cf4d6c7100", + "metadata": {}, + "source": [ + "# Custom Spaces/Kernels\n", + "\n", + "GeometricKernels ships with general-purpose kernel classes that cover most use cases, so you will rarely need to write a new kernel from scratch. In practice, you typically implement a new space and then reuse the existing kernels.\n", + "\n", + "This notebook therefore focuses on implementing custom spaces. At the end, we briefly show how to define a custom kernel for advanced cases (e.g. a non-Matérn kernel or a Matérn kernel with space‑specific optimisation). For an overview of the built-in kernels, see the documentation: https://geometric-kernels.github.io/GeometricKernels/autoapi/geometric_kernels/kernels/index.html.\n", + "\n", + "\n", + "When to implement what:\n", + "\n", + "- **Custom space**: you have a new geometry and can provide the appropriate machinery (e.g. evaluating eignpairs of the Laplacian) for it.\n", + "- **Custom kernel**: you need a non-Matérn kernel or a more efficient implementation tailored to your space." + ] + }, + { + "cell_type": "markdown", + "id": "a041d947-3cc4-44df-86b0-2635dd111d8c", + "metadata": {}, + "source": [ + "## Spaces\n", + "We recommend implementing a new GeometricKernels space by subclassing one of the existing abstract [space classes](https://geometric-kernels.github.io/GeometricKernels/autoapi/geometric_kernels/spaces/base/index.html):\n", + "\n", + "- `DiscreteSpectrumSpace`\n", + "- `NoncompactSymmetricSpace`\n", + "- `HodgeDiscreteSpectrumSpace`\n", + "\n", + "Provide concrete implementations for all abstract methods required by the chosen base class. Doing so will make one of the standard kernel classes directly applicable to your new space:\n", + "\n", + "- `MaternKarhunenLoeveKernel` (for `DiscreteSpectrumSpace`)\n", + "- `MaternFeatureMapKernel` (for `NoncompactSymmetricSpace`)\n", + "- `MaternHodgeCompositionalKernel` (for `HodgeDiscreteSpectrumSpace`)\n", + "\n", + "If you do not want to use the existing kernel machinery, you can instead subclass the most general `Space` class and implement a new kernel from scratch by subclassing `BaseGeometricKernel`.\n", + "\n", + "Notes:\n", + "\n", + "- **Product spaces**: If your target space is a Cartesian product of two or more spaces, you typically do not need to implement a new space class. Instead, consider `ProductGeometricKernel`/`ProductDiscreteSpectrumSpace`. See the documentation on product spaces: https://geometric-kernels.github.io/GeometricKernels/theory/product_spaces.html\n", + "- **Specialising `MaternGeometricKernel`**: If you want to change how `MaternGeometricKernel` specialises to your space (e.g. default approximation level/`num_levels`), consult that class’s documentation and pass the relevant arguments explicitly when constructing the kernel, or change the `kernels.matern_kernel` module.\n", + "- **Contributing a new space**: If you plan to submit a PR adding a new space—thank you! Please integrate it into the test suite and add space-specific tests. Run `make lint` and `make test` before opening a PR. GitHub CI will run these as well (on multiple platforms) and block merging until they pass, but running them locally first is faster." + ] + }, + { + "cell_type": "markdown", + "id": "641d9977-3d89-4061-b07a-339e7b80ea6f", + "metadata": {}, + "source": [ + "**As a toy example**, we implement a space representing a complete graph with $n$ nodes. In practice, this space could be represented using the existing general `Graph` space by providing it the adjacency matrix of the complete graph. However, here we build this space from scratch as if the general `Graph` space did not exist, to demonstrate the mechanics of implementing a new space.\n", + "\n", + "Note: we do not strictly follow the project’s linting/style guidelines in this notebook for a more focused exposition." + ] + }, + { + "cell_type": "markdown", + "id": "4c929868-8bfb-44f8-a1d4-6df376342119", + "metadata": {}, + "source": [ + "### Imports and writing backend‑agnostic code\n", + "\n", + "GeometricKernels works with multiple backends: `pytorch`, `jax`, `numpy`, `tensorflow`. To write code once and have it work across these backends, we use the [lab](https://github.com/wesselb/lab) library. `lab` provides a thin dispatch layer: when you give it, say, PyTorch tensors as inputs, it returns PyTorch tensors; when you give it NumPy arrays, it returns NumPy arrays, etc.\n", + "\n", + "In the next cell we import `lab`." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "3df4b524-381f-44c1-863a-18c0cf4aa420", + "metadata": {}, + "outputs": [], + "source": [ + "import lab as B" + ] + }, + { + "cell_type": "markdown", + "id": "1ec0688c-72db-4560-bab6-9201ee33a844", + "metadata": {}, + "source": [ + "Sometimes, `lab` on its own does not provide all the functions you need. In that case, check the `lab_extras` subpackage in GeometricKernels, which contains additional helpers that follow the same backend-dispatch pattern. If you still cannot find a suitable function, consider adding it to `lab_extras`.\n", + "\n", + "Below we import some utilities we will use:\n", + "- `dtype_integer`: Returns a backend-appropriate integer dtype." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "fe15c4c9-f548-4406-ba20-cca5b354d06b", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "INFO (geometric_kernels): Numpy backend is enabled. To enable other backends, don't forget to `import geometric_kernels.*backend name*`.\n", + "INFO (geometric_kernels): We may be suppressing some logging of external libraries. To override the logging policy, call `logging.basicConfig`.\n" + ] + } + ], + "source": [ + "from geometric_kernels.lab_extras import (\n", + " dtype_integer,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "ba0a00d0-83c0-4bbd-b96d-0064098d0b44", + "metadata": {}, + "source": [ + "Next, we import the abstract base class `DiscreteSpectrumSpace`, because the Laplacian on a finite graph (the graph Laplacian) has a discrete spectrum. See more about different kinds of spaces in the theory docs: https://geometric-kernels.github.io/GeometricKernels/theory/index.html." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "3938680a-ca6a-49c7-8bee-c1f41fbc3975", + "metadata": {}, + "outputs": [], + "source": [ + "from geometric_kernels.spaces.base import DiscreteSpectrumSpace" + ] + }, + { + "cell_type": "markdown", + "id": "640ef02b-41b7-4366-a8af-7259098f089b", + "metadata": {}, + "source": [ + "We also import `Eigenfunctions` because one of the abstract methods of `DiscreteSpectrumSpace` must return an `Eigenfunctions` object.\n", + "\n", + "In many cases, you will subclass `Eigenfunctions` to provide the machinery for working with the Laplacian’s eigenfunctions. In our case, we can avoid that by using the already implemented subclass `EigenfunctionsFromEigenvectors`, which wraps a matrix of eigenvectors." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "77b5f04e-cf3e-4a4b-851b-83304472ce58", + "metadata": {}, + "outputs": [], + "source": [ + "from geometric_kernels.spaces.eigenfunctions import (\n", + " Eigenfunctions,\n", + " EigenfunctionsFromEigenvectors,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "8b120081-988d-4abd-aa29-548d7f783185", + "metadata": {}, + "source": [ + "Finally, a few standard imports. We'll use SciPy to compute the full eigendecomposition for this small example." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "45cb5f35-e473-4284-bedb-e3d775c233a0", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import scipy.linalg" + ] + }, + { + "cell_type": "markdown", + "id": "e60f24cc-bb38-406f-9512-4471284cb137", + "metadata": {}, + "source": [ + "### Defining the new Space Class\n", + "\n", + "Because it is awkward to split a single class definition across multiple Jupyter cells, the complete class definition is given in the next cell.\n", + "\n", + "Note on contributing: If you intend to contribute your new space to GeometricKernels via a PR, please add proper docstrings in the project’s style (these get compiled into the documentation). You can use the implementation of the `Graph` space as a reference. To keep this notebook concise, we include explanatory comments rather than full docstrings." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "9ddb8b66-b547-4556-8667-158dac3ebbdd", + "metadata": {}, + "outputs": [], + "source": [ + "class CompleteGraph(DiscreteSpectrumSpace):\n", + "\n", + " def __init__(self, num_vertices: int):\n", + " # Store the number of vertices.\n", + " self.num_vertices = num_vertices\n", + "\n", + " # Build the adjacency matrix A for the complete graph K_n:\n", + " # A_ij = 1 for i != j, and 0 on the diagonal.\n", + " adjacency_matrix = np.ones((num_vertices, num_vertices), dtype=np.float64) - np.eye(num_vertices, dtype=np.float64)\n", + "\n", + " # Degree matrix D for K_n is (n-1) on the diagonal.\n", + " degree_matrix = np.diag(num_vertices * [num_vertices - 1])\n", + "\n", + " # Normalized graph Laplacian: L = D^(-1/2) (D - A) D^(-1/2).\n", + " self._laplacian = (degree_matrix - adjacency_matrix)/(num_vertices - 1)\n", + "\n", + " # Note: As a rule of thumb, store precomputed quantities in NumPy and cast to the\n", + " # active backend only when needed (usually, kernels handle this casting for you).\n", + "\n", + " # For this toy example, we compute a full eigendecomposition here. In production code,\n", + " # try to keep __init__ lightweight: compute lazily if possible.\n", + " evals, evecs = scipy.linalg.eigh(self._laplacian)\n", + " self._evals = evals\n", + "\n", + " # Re-scale eigenvectors so that they are orthonormal with respect to the discrete inner\n", + " # product induced by the graph (here, multiply by sqrt(n)). This is a convenient normalisation\n", + " # for use in Mercer expansions.\n", + " self._evecs = evecs * num_vertices ** 0.5 # change normalization\n", + "\n", + " def __str__(self):\n", + " # Provide a concise, informative string representation.\n", + " return f\"CompleteGraph({self.num_vertices})\"\n", + "\n", + " @property\n", + " def dimension(self) -> int:\n", + " # Return an intrinsic dimension for the space. For graphs, a pragmatic choice is 0.\n", + " # The dimension influences the parameterisation of the Matérn smoothness parameter nu.\n", + " return 0\n", + "\n", + " def get_eigenfunctions(self, num: int) -> Eigenfunctions:\n", + " # Return the `Eigenfunctions` object representing the first `num` levels of eigenfunctions.\n", + " # The notion of \"levels\" lets many spaces trade off computation for approximation quality.\n", + " # For this space, any admissible `num` has the same one-off cost to precompute, though\n", + " # subsequent kernel evaluations can be faster for smaller `num`.\n", + " #\n", + " # See more about levels in the docs for the `Eigenfunctions` class.\n", + " return EigenfunctionsFromEigenvectors(self._evecs[:, :num])\n", + "\n", + " def get_eigenvalues(self, num: int) -> B.Numeric:\n", + " # Return the first `num` eigenvalues.\n", + " return self._evals[:num]\n", + "\n", + " def get_repeated_eigenvalues(self, num: int) -> B.Numeric:\n", + " # Some DiscreteSpectrumSpace instances group eigenfunctions with multiplicities into \"levels\".\n", + " # In that case, this should return each level’s eigenvalue repeated according to its multiplicity.\n", + " # Here each level corresponds to a single eigenfunction, so we can just return the first `num`\n", + " # values, exactly as in `get_eigenvalues` above.\n", + " #\n", + " # See more about levels in the docs for the `Eigenfunctions` class.\n", + " return self.get_eigenvalues(num)\n", + "\n", + " def random(self, key, number):\n", + " # Generate `number` random vertices (indices) from {0, ..., n-1} in a backend-aware way.\n", + " # This method consumes a backend-specific random key (e.g. a torch.Generator or JAX PRNGKey)\n", + " # and returns samples and the updated key in the same backend.\n", + " #\n", + " # Note this is the first function that actually uses `lab`, the rest delegates supporting\n", + " # different backends to kernels.\n", + " key, random_vertices = B.randint(\n", + " key, dtype_integer(key), number, 1, lower=0, upper=self.num_vertices\n", + " )\n", + " return key, random_vertices\n", + "\n", + " @property\n", + " def element_shape(self):\n", + " # Each point of this space is represented by a single integer node index, thus shape is [1].\n", + " # For R^d you might return [d]; for matrix groups, [d, d], etc. This is used downstream\n", + " # (e.g. for constructing product spaces/kernels).\n", + " return [1]\n", + "\n", + " @property\n", + " def element_dtype(self):\n", + " # The dtype of elements in this space is integer. Also used downstream.\n", + " return B.Int" + ] + }, + { + "cell_type": "markdown", + "id": "13ba08bf-9403-444c-adfa-d780882866bb", + "metadata": {}, + "source": [ + "### Using the new Space\n", + "\n", + "We now demonstrate the space under the `torch` backend. For this, we enable the backend and import a convenience wrapper kernel `MaternGeometricKernel`." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "7e93f56f-561f-4949-b430-aeceb798eb50", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "INFO (geometric_kernels): Torch backend enabled.\n" + ] + } + ], + "source": [ + "import torch\n", + "\n", + "import geometric_kernels\n", + "import geometric_kernels.torch\n", + "\n", + "from geometric_kernels.kernels import MaternGeometricKernel" + ] + }, + { + "cell_type": "markdown", + "id": "1a883a32-c19b-499c-8da1-575dd13d412a", + "metadata": {}, + "source": [ + "Because we subclassed a core abstract space (`DiscreteSpectrumSpace`), we can instantiate `MaternGeometricKernel` on our new space without any additional work." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "24b1cd5d-4ce3-4b0f-b1d0-220644e11846", + "metadata": {}, + "outputs": [], + "source": [ + "mygraph = CompleteGraph(5)\n", + "kernel = MaternGeometricKernel(mygraph)" + ] + }, + { + "cell_type": "markdown", + "id": "1fbf5087-5791-4a37-acbf-f02d50024d0e", + "metadata": {}, + "source": [ + "Initialise the kernel’s parameters. The defaults will come out as NumPy arrays unless you override them." + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "be491823-baa8-4e8d-ac94-b6dae2dec8d6", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "params: {'nu': array([inf]), 'lengthscale': array([1.])}\n" + ] + } + ], + "source": [ + "params = kernel.init_params()\n", + "print(\"params:\", params)" + ] + }, + { + "cell_type": "markdown", + "id": "dee4b49f-b6ff-432c-bbe9-153755f4d550", + "metadata": {}, + "source": [ + "To have the kernel operate in PyTorch, provide parameters as PyTorch tensors; the kernel will then dispatch calculations to the torch backend." + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "0dad5182-6b63-4d3f-b116-0ebe381aa8e1", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "params: {'nu': tensor([1.5000]), 'lengthscale': tensor([2.])}\n" + ] + } + ], + "source": [ + "params[\"lengthscale\"] = torch.tensor([2.0])\n", + "params[\"nu\"] = torch.tensor([3/2])\n", + "print(\"params:\", params)" + ] + }, + { + "cell_type": "markdown", + "id": "c9ebd3a5-c3d5-4b1a-b8b1-2942ec057d1a", + "metadata": {}, + "source": [ + "Generate random inputs for the kernel, also in `torch`. The `random` method consumes and returns the backend’s PRNG key (here `torch.Generator`)." + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "d3d6ef4e-742d-4929-9630-f5b82e4e4817", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "xs: tensor([[0],\n", + " [1],\n", + " [1],\n", + " [0],\n", + " [1],\n", + " [4]], dtype=torch.int32)\n" + ] + } + ], + "source": [ + "key = torch.Generator()\n", + "key.manual_seed(1234)\n", + "\n", + "key, xs = mygraph.random(key, 6)\n", + "\n", + "print(\"xs:\", xs)" + ] + }, + { + "cell_type": "markdown", + "id": "f91f65d0-7596-40ea-a4b4-782f01f730c5", + "metadata": {}, + "source": [ + "Now evaluate the kernel." + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "e4657840-10c0-485f-8d31-51a8843919f8", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "K: tensor([[0.2000, 0.0803, 0.0803, 0.2000, 0.0803, 0.0803],\n", + " [0.0803, 0.2000, 0.2000, 0.0803, 0.2000, 0.0803],\n", + " [0.0803, 0.2000, 0.2000, 0.0803, 0.2000, 0.0803],\n", + " [0.2000, 0.0803, 0.0803, 0.2000, 0.0803, 0.0803],\n", + " [0.0803, 0.2000, 0.2000, 0.0803, 0.2000, 0.0803],\n", + " [0.0803, 0.0803, 0.0803, 0.0803, 0.0803, 0.2000]])\n" + ] + } + ], + "source": [ + "print(\"K:\", kernel.K(params, xs, xs))\n", + "# Note: By default the kernel is normalised such that the variances sum to 1 across nodes.\n", + "# If you prefer marginal variances around 1 for each node, multiply by `num_vertices`." + ] + }, + { + "cell_type": "markdown", + "id": "2b19684f-dcc1-4b16-9d69-6c9743510c76", + "metadata": {}, + "source": [ + "## Kernels\n", + "\n", + "The existing kernel classes cover most use cases; see the [kernel docs](https://geometric-kernels.github.io/GeometricKernels/autoapi/geometric_kernels/kernels/index.html) before rolling your own.\n", + "\n", + "Two common scenarios for a custom kernel are:\n", + "\n", + "1) You want a different PSD kernel that is still representable via the Laplacian’s spectrum (i.e. you only change the spectral filter from the Matérn/heat form).\n", + "2) You want to reimplement the Matérn kernel for your space more efficiently by exploiting finer structure (e.g. you know a closed form expression for your space).\n" + ] + }, + { + "cell_type": "markdown", + "id": "324658e5-7f13-4a76-a16d-b18c8462b080", + "metadata": {}, + "source": [ + "### Custom Kernel for Usecase 1" + ] + }, + { + "cell_type": "markdown", + "id": "a365246e-6c76-4bdf-8de3-05af133cc4b6", + "metadata": {}, + "source": [ + "Below is an example of scenario (1). You can inherit from a default spectral kernel (e.g. `MaternKarhunenLoeveKernel` for graphs) and override the `_spectrum` method, which maps eigenvalues to the Mercer coefficients of the kernel. Here we define an *inverse cosine kernel* on graphs (`Graph` space)." + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "8a89ce18-3539-4b43-ab07-9b815c15d695", + "metadata": {}, + "outputs": [], + "source": [ + "from geometric_kernels.kernels import MaternKarhunenLoeveKernel\n", + "from geometric_kernels.spaces import Graph" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "id": "a7cd30ed-936d-43eb-9616-7410c47c82be", + "metadata": {}, + "outputs": [], + "source": [ + "class InverseCosineKernel(MaternKarhunenLoeveKernel):\n", + " def __init__(self, space: Graph, num_levels: int, normalize: bool = True):\n", + " # For graphs, the Karhunen–Loève-based implementation uses num_levels equal to\n", + " # the number of vertices by default; pass a smaller number to truncate.\n", + " super().__init__(space, num_levels, normalize)\n", + "\n", + " @staticmethod\n", + " def spectrum(\n", + " s: B.Numeric, nu: B.Numeric, lengthscale: B.Numeric, dimension: int\n", + " ) -> B.Numeric:\n", + " # Map eigenvalues s to Mercer coefficients. Ensure this stays non-negative for PSD.\n", + " # **This is only so for the normalized graph Laplacian**.\n", + " # Note: we simply ignore `nu`, `lengthscale` and `dimension` here.\n", + " return B.cos(B.cast(B.dtype(lengthscale), s) * np.pi / 4)" + ] + }, + { + "cell_type": "markdown", + "id": "8792a523-9ff4-4bad-8ba9-c4542b1f1537", + "metadata": {}, + "source": [ + "#### Using the new Kernel" + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "id": "0a992f07-fc92-4346-a0dd-1b14bf8e787d", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "K: tensor([[0.2000, 0.0276, 0.0276, 0.2000, 0.0276, 0.0276],\n", + " [0.0276, 0.2000, 0.2000, 0.0276, 0.2000, 0.0276],\n", + " [0.0276, 0.2000, 0.2000, 0.0276, 0.2000, 0.0276],\n", + " [0.2000, 0.0276, 0.0276, 0.2000, 0.0276, 0.0276],\n", + " [0.0276, 0.2000, 0.2000, 0.0276, 0.2000, 0.0276],\n", + " [0.0276, 0.0276, 0.0276, 0.0276, 0.0276, 0.2000]])\n" + ] + } + ], + "source": [ + "inv_cos_kernel = InverseCosineKernel(mygraph, num_levels=mygraph.num_vertices)\n", + "print(\"K:\", inv_cos_kernel.K(params, xs, xs))" + ] + }, + { + "cell_type": "markdown", + "id": "51009782-6245-4978-939e-9387155c9f3a", + "metadata": {}, + "source": [ + "### Custom Kernel for Usecase 2" + ] + }, + { + "cell_type": "markdown", + "id": "6d326586-b629-4b79-8ff2-7b2a9d169beb", + "metadata": {}, + "source": [ + "For scenario (2), you can inherit from `BaseGeometricKernel` or from a default kernel (e.g. `MaternKarhunenLoeveKernel`) and implement `K` and `K_diag` directly for efficiency.\n", + "\n", + "Below we show how one could, for example, leverage the closed for expression for the heat kernel (Matérn with $\\nu=\\infty$) known for the complete graph." + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "id": "ed139bfe-b56f-4838-9ade-68db7d454d25", + "metadata": {}, + "outputs": [], + "source": [ + "from beartype.typing import Dict, Optional\n", + "\n", + "class MaternKernelCompleteGraph(MaternKarhunenLoeveKernel):\n", + " def __init__(\n", + " self,\n", + " space: CompleteGraph,\n", + " num_levels: int,\n", + " normalize: bool = True,\n", + " ):\n", + " # Call the parent initialiser; override K/K_diag behaviour here if needed.\n", + " super().__init__(space, num_levels, normalize)\n", + "\n", + " def K(\n", + " self,\n", + " params: Dict[str, Dict[str, B.NPNumeric]],\n", + " X: B.Numeric,\n", + " X2: Optional[B.Numeric] = None,\n", + " **kwargs,\n", + " ) -> B.Numeric:\n", + " # Use the exact heat kernel for nu = inf; otherwise defer to the parent. Assumes a normalized Laplacian.\n", + " #\n", + " # Derivation outline for the exact heat kernel (Matérn with nu=inf) on the complete graph.\n", + " # The form is standard (see, e.g., Kondor and Lafferty, 2002), with adjustments for\n", + " # GeometricKernels’ parameterization:\n", + " #\n", + " # Normalized Laplacian for K_n: L = I - D^{-1/2} A D^{-1/2} = (n/(n-1)) I - (1/(n-1)) J, where J is the all-ones matrix.\n", + " # Spectrum: 0 (mult. 1, eigenvector (1, 1, ...)) and α := n/(n-1) (mult. n-1).\n", + " # Heat kernel with length scale κ is exp(- (κ^2 / 2) L).\n", + " # Let t := κ^2 / 2. Then exp(-t L) = J + e^{-α t} (nI - J).\n", + " # Hence entries: k(x, x) = 1 + (n - 1) e^{-α t}, k(x, x') = 1 - e^{-α t} for x != x'.\n", + " nu = params[\"nu\"]\n", + " if nu < np.inf: # Note: this may cause problems with JAX autodiff (see how `MaternKarhunenLoeveKernel` handles conditions)\n", + " return super().K(params, X, X2, **kwargs)\n", + "\n", + " if X2 is None:\n", + " X2 = X\n", + "\n", + " kappa = params[\"lengthscale\"]\n", + " t = kappa**2 / 2\n", + " n = self.space.num_vertices\n", + " alpha = n/(n-1)\n", + "\n", + " k_same = 1 + (n - 1)*B.exp(-alpha*t)\n", + " k_diff = 1 - B.exp(-alpha*t)\n", + " eq = B.cast(B.dtype(kappa), B.eq(X, X2[None, :, 0])) # indicator [x == x']\n", + "\n", + " K = eq*k_same + (1-eq)*k_diff\n", + "\n", + " if self.normalize:\n", + " # Scale so that trace(K) == 1 (sum of marginal variances is 1).\n", + " K = K / k_same * (1/n)\n", + "\n", + " return K\n", + "\n", + " def K_diag(\n", + " self,\n", + " params: Dict[str, Dict[str, B.NPNumeric]],\n", + " X: B.Numeric,\n", + " X2: Optional[B.Numeric] = None,\n", + " **kwargs,\n", + " ) -> B.Numeric:\n", + " # Normally, you would implement this separately to reduce complexity.\n", + " return self.K(params, X, X2, **kwargs)" + ] + }, + { + "cell_type": "markdown", + "id": "a164d068-d7c7-4819-8e81-b64d0d21d5d4", + "metadata": {}, + "source": [ + "#### Checking the new Kernel\n", + "First, $\\nu < \\infty$, effectively using the default kernel." + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "id": "a1f7ab13-4424-45fd-905b-b433c9acc170", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "K_alt:\n", + " tensor([[0.2000, 0.0803, 0.0803, 0.2000, 0.0803, 0.0803],\n", + " [0.0803, 0.2000, 0.2000, 0.0803, 0.2000, 0.0803],\n", + " [0.0803, 0.2000, 0.2000, 0.0803, 0.2000, 0.0803],\n", + " [0.2000, 0.0803, 0.0803, 0.2000, 0.0803, 0.0803],\n", + " [0.0803, 0.2000, 0.2000, 0.0803, 0.2000, 0.0803],\n", + " [0.0803, 0.0803, 0.0803, 0.0803, 0.0803, 0.2000]])\n", + "K_orig:\n", + " tensor([[0.2000, 0.0803, 0.0803, 0.2000, 0.0803, 0.0803],\n", + " [0.0803, 0.2000, 0.2000, 0.0803, 0.2000, 0.0803],\n", + " [0.0803, 0.2000, 0.2000, 0.0803, 0.2000, 0.0803],\n", + " [0.2000, 0.0803, 0.0803, 0.2000, 0.0803, 0.0803],\n", + " [0.0803, 0.2000, 0.2000, 0.0803, 0.2000, 0.0803],\n", + " [0.0803, 0.0803, 0.0803, 0.0803, 0.0803, 0.2000]])\n" + ] + } + ], + "source": [ + "alt_kernel = MaternKernelCompleteGraph(mygraph, mygraph.num_vertices)\n", + "\n", + "print(\"K_alt:\\n\", alt_kernel.K(params, xs, xs))\n", + "print(\"K_orig:\\n\", kernel.K(params, xs, xs))" + ] + }, + { + "cell_type": "markdown", + "id": "3fd7ca1b-981a-478c-80c6-c9daabf89a2f", + "metadata": {}, + "source": [ + "Next, $\\nu = \\infty$, using the new code." + ] + }, + { + "cell_type": "code", + "execution_count": 18, + "id": "2cff7f2d-70df-4ff2-8408-21ae3217026a", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "K_alt:\n", + " tensor([[0.2000, 0.1382, 0.1382, 0.2000, 0.1382, 0.1382],\n", + " [0.1382, 0.2000, 0.2000, 0.1382, 0.2000, 0.1382],\n", + " [0.1382, 0.2000, 0.2000, 0.1382, 0.2000, 0.1382],\n", + " [0.2000, 0.1382, 0.1382, 0.2000, 0.1382, 0.1382],\n", + " [0.1382, 0.2000, 0.2000, 0.1382, 0.2000, 0.1382],\n", + " [0.1382, 0.1382, 0.1382, 0.1382, 0.1382, 0.2000]])\n", + "K_orig:\n", + " tensor([[0.2000, 0.1382, 0.1382, 0.2000, 0.1382, 0.1382],\n", + " [0.1382, 0.2000, 0.2000, 0.1382, 0.2000, 0.1382],\n", + " [0.1382, 0.2000, 0.2000, 0.1382, 0.2000, 0.1382],\n", + " [0.2000, 0.1382, 0.1382, 0.2000, 0.1382, 0.1382],\n", + " [0.1382, 0.2000, 0.2000, 0.1382, 0.2000, 0.1382],\n", + " [0.1382, 0.1382, 0.1382, 0.1382, 0.1382, 0.2000]])\n" + ] + } + ], + "source": [ + "alt_params = params.copy()\n", + "alt_params[\"nu\"] = torch.tensor([np.inf])\n", + "\n", + "print(\"K_alt:\\n\", alt_kernel.K(alt_params, xs, xs))\n", + "print(\"K_orig:\\n\", kernel.K(alt_params, xs, xs))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5ab88ed9-7123-4825-b43c-1e1d3e4a47fd", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "gkconda_updated_sh", + "language": "python", + "name": "gkconda_updated_sh" + }, + "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.10.13" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}