Merge branch 'main' into slava/python312

CHANGELOG.md

@@ -1,5 +1,17 @@
 # CHANGELOG
 
+## v0.2.1 - 08.08.2024
+ Minor release with mostly cosmetic changes:
+* Add "If you have a question" section to README.md by @vabor112 in https://github.com/geometric-kernels/GeometricKernels/pull/131
+* Github cosmetics by @stoprightthere in https://github.com/geometric-kernels/GeometricKernels/pull/133
+* Replace all references to "gpflow" organization with "geometric-kernels" organization by @vabor112 in https://github.com/geometric-kernels/GeometricKernels/pull/134
+* Use fit_gpytorch_model or fit.fit_gpytorch_mll depening on the botorсh version by @vabor112 in https://github.com/geometric-kernels/GeometricKernels/pull/137
+* Add a missing type cast and fix a typo in kernels/karhunen_loeve.py by @vabor112 in https://github.com/geometric-kernels/GeometricKernels/pull/136
+* Minor documentation improvements by @vabor112 in https://github.com/geometric-kernels/GeometricKernels/pull/135
+* Add citation to the preprint of the GeometricKernels paper by @vabor112 in https://github.com/geometric-kernels/GeometricKernels/pull/138
+* Add citation file by @aterenin in https://github.com/geometric-kernels/GeometricKernels/pull/140
+* Fix dependencies (Version 0.2.1) by @stoprightthere in https://github.com/geometric-kernels/GeometricKernels/pull/143
+
 ## v0.2 - 21.04.2024
  New geometric kernel that *just works*, `kernels.MaternGeometricKernel`. Relies on *(hopefully)* sensible defaults we defined. Mostly by @stoprightthere.
 

CITATION.bib

@@ -0,0 +1,6 @@
+@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},
+}

CITATION.cff

@@ -0,0 +1,29 @@
+cff-version: 1.2.0
+message: "If you use this software, please cite it as below."
+title: "GeometricKernels"
+authors:
+- name: "GeometricKernels Contributors"
+preferred-citation:
+  type: "article"
+  title: "The GeometricKernels Package: Heat and Matérn Kernels for Geometric Learning on Manifolds, Meshes, and Graphs"
+  authors:
+  - family-names: "Mostowsky"
+    given-names: "Peter"
+  - family-names: "Dutordoir"
+    given-names: "Vincent"
+  - family-names: "Azangulov"
+    given-names: "Iskander"
+  - family-names: "Jaquier"
+    given-names: "Noémie"
+  - family-names: "Hutchinson"
+    given-names: "Michael John"
+  - family-names: "Ravuri"
+    given-names: "Aditya"
+  - family-names: "Rozo"
+    given-names: "Leonel"
+  - family-names: "Terenin"
+    given-names: "Alexander"
+  - family-names: "Borovitskiy"
+    given-names: "Viacheslav"
+  year: "2024"
+  journal: "arXiv:2407.08086"

docs/examples/HypercubeGraph.nblink

@@ -0,0 +1 @@
+{"path": "../../notebooks/HypercubeGraph.ipynb"}

docs/examples/examples.rst

@@ -8,10 +8,11 @@ Spaces
 
    Graph <Graph>
    Hyperbolic <Hyperbolic>
+   Hypercube <Hypercube>
    Hypersphere <Hypersphere>
    Mesh <Mesh>
-   SpecialOrthogonal <SpecialOrthogonal>
    SPD <SPD>
+   SpecialOrthogonal <SpecialOrthogonal>
    SpecialUnitary <SpecialUnitary>
    Torus <Torus>
 

docs/references.bib

@@ -121,4 +121,18 @@ @article{canzani2013
   author={Canzani, Yaiza},
   journal={Lecture Notes available at: http://www.math.harvard.edu/canzani/docs/Laplacian.pdf},
   year={2013}
+}
+
+@book{macwilliams1977,
+  title={The theory of error-correcting codes},
+  author={MacWilliams, Florence Jessie and Sloane, Neil James Alexander},
+  year={1977},
+  publisher={Elsevier}
+}
+
+@inproceedings{borovitskiy2023,
+  title={Isotropic Gaussian Processes on Finite Spaces of Graphs},
+  author={Borovitskiy, Viacheslav and Karimi, Mohammad Reza and Somnath, Vignesh Ram and Krause, Andreas},
+  booktitle={International Conference on Artificial Intelligence and Statistics},
+  year={2023},
 }

docs/theory/hypercube_graph.rst

@@ -0,0 +1,108 @@
+################################
+  Kernels on the Hypercube Graph
+################################
+
+.. warning::
+    You can get by fine without reading this page for almost all use cases, just use the standard :class:`~.kernels.MaternGeometricKernel`, following the respective :doc:`example notebook </examples/HypercubeGraph>`.
+
+    This is optional material meant to explain the basic theory and based mainly on :cite:t:`borovitskiy2023`.
+
+==========================
+Motivation
+==========================
+
+The :class:`~.spaces.HypercubeGraph` space $C^d$ can be used to model $d$-dimensional *binary vector* inputs.
+
+There are many settings where inputs are binary vectors or can be represented as such. For instance, upon flattening, binary vectors represent adjacency matrices of *unweighted labeled graphs* [#]_. 
+
+==========================
+Structure of the Space
+==========================
+
+The elements of this space—given its `dim` is $d \in \mathbb{Z}_{>0}$—are exactly the binary vectors of length $d$.
+
+The geometry of this space is simple: it is a graph such that $x, x' \in C^d$ are connected by an edge if and only if they differ in exactly one coordinate (i.e. there is exactly *Hamming distance* $1$ between).
+
+Being a graph, $C^d$ could also be represented using the general :class:`~.spaces.Graph` space.
+However, the number of nodes in $C^d$ is $2^d$, which is exponential in $d$, rendering the general techniques infeasible.
+
+==========================
+Eigenfunctions
+==========================
+
+On graphs, kernels are computed using the eigenfunctions and eigenvalues of the Laplacian.
+
+The eigenfunctions of the Laplacian on the hypercube graph are the *Walsh functions* [#]_ given analytically by the simple formula
+$$
+w_T(x_0, .., x_{d-1}) = (-1)^{\sum_{j \in T} x_j}
+$$
+where $x = (x_0, .., x_{d-1}) \in C^d$ and the index $T$ is an arbitrary subset of the set $\{0, .., d-1\}$.
+
+The corresponding eigenvalues are $\lambda_T = \lambda_{\lvert T \rvert} = 2 \lvert T \rvert / d$, where $\lvert T \rvert$ is the cardinality of $T$.
+
+However, the problem is that the number of eigenfunctions is $2^d$.
+Hence naive truncation of the sum in the kernel formula to a few hundred terms leads to a poor approximation of the kernel for larger $d$.
+
+==========================
+Addition Theorem
+==========================
+
+Much like for the hyperspheres and unlike for the general graphs, there is an :doc:`addition theorem <addition_theorem>` for the hypercube graph:
+
+$$
+\sum_{T \subseteq \{0, .., d-1\}, \lvert T \rvert = j} w_T(x) w_T(x')
+=
+\sum_{T \subseteq \{0, .., d-1\}, \lvert T \rvert = j} w_T(x \oplus x')
+=
+\binom{d}{j}
+\widetilde{G}_{d, j, m}
+$$
+where $\oplus$ is the elementwise XOR operation, $m$ is the Hamming distance between $x$ and $x'$, and $\widetilde{G}_{d, j, m}$ is the Kravchuk polynomial of degree $d$ and order $j$ normalized such that $\widetilde{G}_{d, j, 0} = 1$, evaluated at $m$.
+
+Normalized Kravchuk polynomials $\widetilde{G}_{d, j, m}$ satisfy the following three-term recurrence relation
+$$
+\widetilde{G}_{d, j, m}
+=
+\frac{d - 2 m}{d - j + 1} \widetilde{G}_{d, j - 1, m}
+-\frac{j-1}{d - j + 1} \widetilde{G}_{d, j - 2, m},
+\quad
+\widetilde{G}_{d, 0, m} = 1,
+\quad
+\widetilde{G}_{d, 1, m} = 1 - \frac{2}{d} m,
+$$
+which allows for their efficient computation without the need to compute large sums of the individual Walsh functions.
+
+With that, the kernels on the hypercube graph can be computed efficiently using the formula
+$$
+k_{\nu, \kappa}(x, x')
+=
+\frac{1}{C_{\nu, \kappa}}
+\sum_{l=0}^{L-1}
+\Phi_{\nu, \kappa}(\lambda_l)
+\binom{d}{j} \widetilde{G}_{d, j, m}
+\qquad
+\Phi_{\nu, \kappa}(\lambda)
+=
+\begin{cases}
+\left(\frac{2\nu}{\kappa^2} + \lambda\right)^{-\nu-\frac{d}{2}}
+&
+\nu < \infty \text{ — Matérn}
+\\
+e^{-\frac{\kappa^2}{2} \lambda}
+&
+\nu = \infty \text{ — Heat (RBF)}
+\end{cases}
+$$
+where $m$ is the Hamming distance between $x$ and $x'$, and $L \leq d + 1$ is the user-controlled number of levels parameters.
+
+**Notes:**
+
+#. We define the dimension of the :class:`~.spaces.HypercubeGraph` space $C^d$ to be $d$, in contrast to the graphs represented by the :class:`~.spaces.Graph` space, whose dimension is defined to be $0$.
+
+   Because of this, much like in the Euclidean or the manifold case, the $1/2, 3/2, 5/2$ *are* in fact reasonable values of for the smoothness parameter $\nu$.
+
+.. rubric:: Footnotes
+
+.. [#] Every node of a labeled graph is associated with a unique label. Functions on labeled graphs do *not* have to be invariant to permutations of nodes.
+
+.. [#] Since the hypercube graph $C^d$ is $d$-regular, the unnormalized Laplacian and the symmetric normalized Laplacian coincide up to a multiplication by $d$. Thus their eigenfunctions are the same and eigenvalues coincide up to a multiplication by $d$. For better numerical stability, we use symmetric normalized Laplacian in the implementation and assume its use throughout this page.

docs/theory/index.rst

@@ -13,3 +13,4 @@
    Kernels on product spaces <product_spaces>
    Product kernels <product_kernels>
    Feature maps and sampling <feature_maps>
+   Hypercube graph space <hypercube_graph>

geometric_kernels/feature_maps/probability_densities.py

@@ -13,7 +13,6 @@
 import lab as B
 import numpy as np
 from beartype.typing import Dict, List, Optional, Tuple
-from opt_einsum import contract as einsum
 from sympy import Poly, Product, symbols
 
 from geometric_kernels.lab_extras import (
@@ -85,7 +84,7 @@ def student_t_sample(
     shape_sqrt = B.chol(shape)
     dtype = dtype or dtype_double(key)
     key, z = B.randn(key, dtype, *size, n)
-    z = einsum("...i,ji->...j", z, shape_sqrt)
+    z = B.einsum("...i,ji->...j", z, shape_sqrt)
 
     key, g = B.randgamma(
         key,

geometric_kernels/kernels/feature_map.py

@@ -7,7 +7,6 @@
 import lab as B
 import numpy as np
 from beartype.typing import Dict, Optional
-from opt_einsum import contract as einsum
 
 from geometric_kernels.feature_maps import FeatureMap
 from geometric_kernels.kernels.base import BaseGeometricKernel
@@ -124,7 +123,7 @@ def K(
         else:
             features_X2 = features_X
 
-        feature_product = einsum("...no,...mo->...nm", features_X, features_X2)
+        feature_product = B.einsum("...no,...mo->...nm", features_X, features_X2)
         return feature_product
 
     def K_diag(self, params: Dict[str, B.Numeric], X: B.Numeric, **kwargs):

geometric_kernels/kernels/matern_kernel.py

@@ -24,6 +24,7 @@
     DiscreteSpectrumSpace,
     Graph,
     Hyperbolic,
+    HypercubeGraph,
     Hypersphere,
     Mesh,
     NoncompactSymmetricSpace,
@@ -199,6 +200,8 @@ def default_num(space: DiscreteSpectrumSpace) -> int:
         return min(
             MaternGeometricKernel._DEFAULT_NUM_EIGENFUNCTIONS, space.num_vertices
         )
+    elif isinstance(space, HypercubeGraph):
+        return min(MaternGeometricKernel._DEFAULT_NUM_LEVELS, space.dim + 1)
     else:
         return MaternGeometricKernel._DEFAULT_NUM_LEVELS
 

geometric_kernels/lab_extras/extras.py

@@ -308,3 +308,38 @@ def complex_conj(x: B.Numeric):
     :param x:
         Array of any backend.
     """
+
+
+@dispatch
+@abstract()
+def logical_xor(x1: B.Bool, x2: B.Bool):
+    """
+    Return logical XOR of two arrays.
+
+    :param x1:
+        Array of any backend.
+    :param x2:
+        Array of any backend.
+    """
+
+
+@dispatch
+@abstract()
+def count_nonzero(x: B.Numeric, axis=None):
+    """
+    Count non-zero elements in an array.
+
+    :param x:
+        Array of any backend and of any shape.
+    """
+
+
+@dispatch
+@abstract()
+def dtype_bool(reference: B.RandomState):
+    """
+    Return `bool` dtype of a backend based on the reference.
+
+    :param reference:
+        A random state to infer the backend from.
+    """

geometric_kernels/lab_extras/jax/extras.py

@@ -221,3 +221,27 @@ def complex_conj(x: B.JAXNumeric):
     Return complex conjugate
     """
     return jnp.conj(x)
+
+
+@dispatch
+def logical_xor(x1: B.JAXNumeric, x2: B.JAXNumeric):
+    """
+    Return logical XOR of two arrays.
+    """
+    return jnp.logical_xor(x1, x2)
+
+
+@dispatch
+def count_nonzero(x: B.JAXNumeric, axis=None):
+    """
+    Count non-zero elements in an array.
+    """
+    return jnp.count_nonzero(x, axis=axis)
+
+
+@dispatch
+def dtype_bool(reference: B.JAXRandomState):  # type: ignore
+    """
+    Return `bool` dtype of a backend based on the reference.
+    """
+    return bool

geometric_kernels/lab_extras/numpy/extras.py

@@ -216,3 +216,27 @@ def complex_conj(x: B.NPNumeric):
     Return complex conjugate
     """
     return np.conjugate(x)
+
+
+@dispatch
+def logical_xor(x1: B.NPNumeric, x2: B.NPNumeric):
+    """
+    Return logical XOR of two arrays.
+    """
+    return np.logical_xor(x1, x2)
+
+
+@dispatch
+def count_nonzero(x: B.NPNumeric, axis=None):
+    """
+    Count non-zero elements in an array.
+    """
+    return np.count_nonzero(x, axis=axis)
+
+
+@dispatch
+def dtype_bool(reference: B.NPRandomState):  # type: ignore
+    """
+    Return `bool` dtype of a backend based on the reference.
+    """
+    return bool

geometric_kernels/lab_extras/tensorflow/extras.py

@@ -227,3 +227,27 @@ def complex_conj(x: B.TFNumeric):
     Return complex conjugate
     """
     return tf.math.conj(x)
+
+
+@dispatch
+def logical_xor(x1: B.TFNumeric, x2: B.TFNumeric):
+    """
+    Return logical XOR of two arrays.
+    """
+    return tf.math.logical_xor(x1, x2)
+
+
+@dispatch
+def count_nonzero(x: B.TFNumeric, axis=None):
+    """
+    Count non-zero elements in an array.
+    """
+    return tf.math.count_nonzero(x, axis=axis)
+
+
+@dispatch
+def dtype_bool(reference: B.TFRandomState):  # type: ignore
+    """
+    Return `bool` dtype of a backend based on the reference.
+    """
+    return tf.bool