diff --git a/Makefile b/Makefile index 76b0decb..9181219c 100644 --- a/Makefile +++ b/Makefile @@ -34,5 +34,5 @@ lint: test: ## Run the tests, start with the failing ones and break on first fail. - pytest -v -x --ff -rN -Wignore -s --tb=short --durations=0 tests + pytest -v -x --ff -rN -Wignore -s --tb=short --durations=0 --cov --cov-report=xml tests pytest --nbmake --nbmake-kernel=python3 --durations=0 --nbmake-timeout=1000 --ignore=notebooks/frontends/GPJax.ipynb notebooks/ diff --git a/docs/references.bib b/docs/references.bib index b33311ad..0328bfa3 100644 --- a/docs/references.bib +++ b/docs/references.bib @@ -46,18 +46,24 @@ @inproceedings{jaquier2021 year={2021} } -@article{azangulov2022, - title = {Stationary Kernels and Gaussian Processes on Lie Groups and their Homogeneous Spaces I: the compact case}, - author = {Azangulov, Iskander and Smolensky, Andrei and Terenin, Alexander and Borovitskiy, Viacheslav}, - journal = {arXiv preprint arXiv:2208.14960}, - year = {2022} +@article{azangulov2024a, + title={Stationary Kernels and Gaussian Processes on Lie Groups and their Homogeneous Spaces I: the compact case}, + author={Azangulov, Iskander and Smolensky, Andrei and Terenin, Alexander and Borovitskiy, Viacheslav}, + journal={Journal of Machine Learning Research}, + year={2024}, + volume={25}, + number={280}, + pages={1--52}, } -@article{azangulov2023, +@article{azangulov2024b, title={Stationary Kernels and Gaussian Processes on Lie Groups and their Homogeneous Spaces II: non-compact symmetric spaces}, author={Azangulov, Iskander and Smolensky, Andrei and Terenin, Alexander and Borovitskiy, Viacheslav}, - journal={arXiv preprint arXiv:2301.13088}, - year={2023} + journal={Journal of Machine Learning Research}, + year={2024}, + volume={25}, + number={281}, + pages={1--51}, } @inproceedings{yang2024, @@ -135,4 +141,15 @@ @inproceedings{borovitskiy2023 author={Borovitskiy, Viacheslav and Karimi, Mohammad Reza and Somnath, Vignesh Ram and Krause, Andreas}, booktitle={International Conference on Artificial Intelligence and Statistics}, year={2023}, +} + +@article{sawyer1992, + author = {Sawyer, Patrice}, + journal = {Canadian Journal of Mathematics}, + number = {3}, + pages = {624--651}, + publisher = {Cambridge University Press}, + title = {The heat equation on the spaces of positive definite matrices}, + volume = {44}, + year = {1992}, } \ No newline at end of file diff --git a/docs/theory/addition_theorem.rst b/docs/theory/addition_theorem.rst index 7fd79ca5..58d58c0b 100644 --- a/docs/theory/addition_theorem.rst +++ b/docs/theory/addition_theorem.rst @@ -68,7 +68,7 @@ For example, for $M = \mathbb{S}_2$ and $L = 20$ the corresponding $J$ is $400$. In the simplest special case of $\mathbb{S}_d$, the circle $\mathbb{S}_1$, the eigenfunctions are given by $\sin(l \theta), \cos(l \theta)$, where $l$ indexes levels. The outer product $\cos(l \theta) \cos(l \theta') + \sin(l \theta) \sin(l \theta')$ in this case can be simplified to $\cos(l (\theta-\theta')) = \cos(l d_{\mathbb{S}_1}(\theta, \theta'))$ thanks to an elementary trigonometric identity. -Such addition theorems appear beyond hyperspheres, for example for Lie groups and other compact homogeneous spaces :cite:p:`azangulov2022`. +Such addition theorems appear beyond hyperspheres, for example for Lie groups and other compact homogeneous spaces :cite:p:`azangulov2024a`. In the library, such spaces use the class :class:`~.EigenfunctionsWithAdditionTheorem` to represent the spectrum of $\Delta_{\mathcal{M}}$. For them, the *number of levels* parameter of the :class:`~.kernels.MaternKarhunenLoeveKernel` maps to $L$ in the above formula. diff --git a/docs/theory/symmetric.rst b/docs/theory/symmetric.rst index 33ae284a..546c305c 100644 --- a/docs/theory/symmetric.rst +++ b/docs/theory/symmetric.rst @@ -5,7 +5,7 @@ .. warning:: You can get by fine without reading this page for almost all use cases, just use the standard :class:`~.kernels.MaternGeometricKernel`, following the example notebooks :doc:`on hyperbolic spaces ` and :doc:`on the space of symmetric positive definite matrices (SPD) `. - This is optional material meant to explain the basic theory and based mainly on :cite:t:`azangulov2023`. + This is optional material meant to explain the basic theory and based mainly on :cite:t:`azangulov2024b`. ======= Theory @@ -20,7 +20,7 @@ In the Euclidean case, closed form expressions for kernels are available and ran No closed form expressions for kernels are usually available on other non-compact symmetric spaces. Because of that, random Fourier features are the basic means of computing the kernels in this case. -A complete mathematical treatise can be found in :cite:t:`azangulov2023`. +A complete mathematical treatise can be found in :cite:t:`azangulov2024b`. Here we briefly present the main ideas. Recall that the usual Euclidean random Fourier features boil down to diff --git a/geometric_kernels/feature_maps/deterministic.py b/geometric_kernels/feature_maps/deterministic.py index 873a57d0..60906706 100644 --- a/geometric_kernels/feature_maps/deterministic.py +++ b/geometric_kernels/feature_maps/deterministic.py @@ -9,7 +9,6 @@ from beartype.typing import Dict, Optional, Tuple from geometric_kernels.feature_maps.base import FeatureMap -from geometric_kernels.kernels.karhunen_loeve import MaternKarhunenLoeveKernel from geometric_kernels.spaces import DiscreteSpectrumSpace @@ -25,6 +24,8 @@ class DeterministicFeatureMapCompact(FeatureMap): """ def __init__(self, space: DiscreteSpectrumSpace, num_levels: int): + from geometric_kernels.kernels.karhunen_loeve import MaternKarhunenLoeveKernel + self.space = space self.num_levels = num_levels self.kernel = MaternKarhunenLoeveKernel(space, num_levels) diff --git a/geometric_kernels/feature_maps/probability_densities.py b/geometric_kernels/feature_maps/probability_densities.py index 51e524cd..1e472ccc 100644 --- a/geometric_kernels/feature_maps/probability_densities.py +++ b/geometric_kernels/feature_maps/probability_densities.py @@ -217,7 +217,7 @@ def _randcat_fix( def _alphas(n: int) -> B.Numeric: r""" - Compute alphas for Prop. 16 & 17 of cite:t:`azangulov2023` + Compute alphas for Prop. 16 & 17 of cite:t:`azangulov2024b` for the hyperbolic space of dimension `n`. :param n: @@ -299,7 +299,7 @@ def _sample_mixture_matern( ) -> Tuple[B.RandomState, B.Numeric]: r""" Sample from the mixture distribution from Prop. 17 of - cite:t:`azangulov2023` for specific alphas `alpha`, length + cite:t:`azangulov2024b` for specific alphas `alpha`, length scale ($\kappa$) `lengthscale`, smoothness `nu` and dimension `dim`, using `key` random state. diff --git a/geometric_kernels/feature_maps/random_phase.py b/geometric_kernels/feature_maps/random_phase.py index 79a55111..cb1d461a 100644 --- a/geometric_kernels/feature_maps/random_phase.py +++ b/geometric_kernels/feature_maps/random_phase.py @@ -17,7 +17,6 @@ from geometric_kernels.feature_maps.base import FeatureMap from geometric_kernels.feature_maps.probability_densities import base_density_sample -from geometric_kernels.kernels.karhunen_loeve import MaternKarhunenLoeveKernel from geometric_kernels.lab_extras import complex_like, from_numpy, is_complex from geometric_kernels.spaces import DiscreteSpectrumSpace, NoncompactSymmetricSpace @@ -43,6 +42,8 @@ def __init__( num_levels: int, num_random_phases: int = 3000, ): + from geometric_kernels.kernels.karhunen_loeve import MaternKarhunenLoeveKernel + self.space = space self.num_levels = num_levels self.num_random_phases = num_random_phases diff --git a/geometric_kernels/kernels/product.py b/geometric_kernels/kernels/product.py index 5f3a576d..69dcdde6 100644 --- a/geometric_kernels/kernels/product.py +++ b/geometric_kernels/kernels/product.py @@ -66,6 +66,7 @@ def __init__( assert isinstance(kernel.space, Space) self.spaces.append(kernel.space) self.element_shapes = [space.element_shape for space in self.spaces] + self.element_dtypes = [space.element_dtype for space in self.spaces] if dimension_indices is None: dimensions = [math.prod(shape) for shape in self.element_shapes] @@ -114,9 +115,13 @@ def K(self, params: Dict[str, B.Numeric], X, X2=None, **kwargs) -> B.Numeric: if X2 is None: X2 = X - Xs = project_product(X, self.dimension_indices, self.element_shapes) - X2s = project_product(X2, self.dimension_indices, self.element_shapes) - params_list = params_to_params_list(params) + Xs = project_product( + X, self.dimension_indices, self.element_shapes, self.element_dtypes + ) + X2s = project_product( + X2, self.dimension_indices, self.element_shapes, self.element_dtypes + ) + params_list = params_to_params_list(len(self.kernels), params) return B.prod( B.stack( @@ -130,8 +135,10 @@ def K(self, params: Dict[str, B.Numeric], X, X2=None, **kwargs) -> B.Numeric: ) def K_diag(self, params, X): - Xs = project_product(X, self.dimension_indices, self.element_shapes) - params_list = params_to_params_list(params) + Xs = project_product( + X, self.dimension_indices, self.element_shapes, self.element_dtypes + ) + params_list = params_to_params_list(len(self.kernels), params) return B.prod( B.stack( diff --git a/geometric_kernels/lab_extras/extras.py b/geometric_kernels/lab_extras/extras.py index f1076d01..8d75a4fd 100644 --- a/geometric_kernels/lab_extras/extras.py +++ b/geometric_kernels/lab_extras/extras.py @@ -343,3 +343,37 @@ def dtype_bool(reference: B.RandomState): :param reference: A random state to infer the backend from. """ + + +@dispatch +@abstract() +def bool_like(reference: B.Numeric): + """ + Return the type of the reference if it is of boolean type. + Otherwise return `bool` dtype of a backend based on the reference. + + :param reference: + Array of any backend. + """ + + +def smart_cast( + dtype: Union[B.Bool, B.Int, B.Float, B.Complex, B.Numeric], x: B.Numeric +): + """ + Return `x` cast to the `dtype` abstract data type. + + :param dtype: + An abstract DType of lab, one of `B.Bool`, `B.Int`, `B.Float`, + `B.Complex`, `B.Numeric`. + :param x: + Array of any backend. + """ + if dtype == B.Bool: + return B.cast(bool_like(x), x) + elif dtype == B.Int: + return B.cast(int_like(x), x) + elif dtype == B.Float: + return B.cast(float_like(x), x) + elif dtype == B.Complex: + return B.cast(complex_like(x), x) diff --git a/geometric_kernels/lab_extras/jax/extras.py b/geometric_kernels/lab_extras/jax/extras.py index 81c93846..5743ef3e 100644 --- a/geometric_kernels/lab_extras/jax/extras.py +++ b/geometric_kernels/lab_extras/jax/extras.py @@ -2,7 +2,7 @@ import lab as B from beartype.typing import List from lab import dispatch -from plum import Union +from plum import Union, convert _Numeric = Union[B.Number, B.JAXNumeric] @@ -89,7 +89,9 @@ def float_like(reference: B.JAXNumeric): """ reference_dtype = reference.dtype if jnp.issubdtype(reference_dtype, jnp.floating): - return reference_dtype + return convert( + reference_dtype, B.JAXDType + ) # JAX .dtype returns a NumPy data type. This converts it to a JAX one. else: return jnp.float64 @@ -106,7 +108,9 @@ def dtype_integer(reference: B.JAXRandomState): # type: ignore def int_like(reference: B.JAXNumeric): reference_dtype = reference.dtype if jnp.issubdtype(reference_dtype, jnp.integer): - return reference_dtype + return convert( + reference_dtype, B.JAXDType + ) # JAX .dtype returns a NumPy data type. This converts it to a JAX one. else: return jnp.int32 @@ -155,10 +159,7 @@ def complex_like(reference: B.JAXNumeric): """ Return `complex` dtype of a backend based on the reference. """ - if B.dtype(reference) == jnp.float32: - return jnp.complex64 - else: - return jnp.complex128 + return B.promote_dtypes(jnp.complex64, reference.dtype) @dispatch @@ -244,4 +245,19 @@ def dtype_bool(reference: B.JAXRandomState): # type: ignore """ Return `bool` dtype of a backend based on the reference. """ - return bool + return jnp.bool_ + + +@dispatch +def bool_like(reference: B.JAXRandomState): + """ + Return the type of the reference if it is of boolean type. + Otherwise return `bool` dtype of a backend based on the reference. + """ + reference_dtype = reference.dtype + if jnp.issubdtype(reference_dtype, jnp.bool_): + return convert( + reference_dtype, B.JAXDType + ) # JAX .dtype returns a NumPy data type. This converts it to a JAX one. + else: + return jnp.bool_ diff --git a/geometric_kernels/lab_extras/numpy/extras.py b/geometric_kernels/lab_extras/numpy/extras.py index 1d9fe038..01c69168 100644 --- a/geometric_kernels/lab_extras/numpy/extras.py +++ b/geometric_kernels/lab_extras/numpy/extras.py @@ -144,10 +144,7 @@ def complex_like(reference: B.NPNumeric): """ Return `complex` dtype of a backend based on the reference. """ - if reference.dtype == np.float32: - return np.complex64 - else: - return np.complex128 + return B.promote_dtypes(np.complex64, reference.dtype) @dispatch @@ -239,4 +236,17 @@ def dtype_bool(reference: B.NPRandomState): # type: ignore """ Return `bool` dtype of a backend based on the reference. """ - return bool + return np.bool_ + + +@dispatch +def bool_like(reference: B.NPNumeric): + """ + Return the type of the reference if it is of boolean type. + Otherwise return `bool` dtype of a backend based on the reference. + """ + reference_dtype = reference.dtype + if np.issubdtype(reference_dtype, np.bool_): + return reference_dtype + else: + return np.bool_ diff --git a/geometric_kernels/lab_extras/numpy/sparse_extras.py b/geometric_kernels/lab_extras/numpy/sparse_extras.py index 19c316be..e4a938ce 100644 --- a/geometric_kernels/lab_extras/numpy/sparse_extras.py +++ b/geometric_kernels/lab_extras/numpy/sparse_extras.py @@ -1,3 +1,5 @@ +import sys + import lab as B import scipy import scipy.sparse as sp @@ -11,15 +13,21 @@ SparseArray defines a lab data type that covers all possible sparse scipy arrays, so that multiple dispatch works with such arrays. """ -SparseArray = Union[ - sp.bsr_matrix, - sp.coo_matrix, - sp.csc_matrix, - sp.csr_matrix, - sp.dia_matrix, - sp.dok_matrix, - sp.lil_matrix, -] +if sys.version_info[:2] <= (3, 8): + SparseArray = Union[ + sp.bsr_matrix, + sp.coo_matrix, + sp.csc_matrix, + sp.csr_matrix, + sp.dia_matrix, + sp.dok_matrix, + sp.lil_matrix, + ] +else: + SparseArray = Union[ + sp.sparray, + sp.spmatrix, + ] @dispatch diff --git a/geometric_kernels/lab_extras/tensorflow/extras.py b/geometric_kernels/lab_extras/tensorflow/extras.py index 8e613f15..594be7f9 100644 --- a/geometric_kernels/lab_extras/tensorflow/extras.py +++ b/geometric_kernels/lab_extras/tensorflow/extras.py @@ -1,3 +1,5 @@ +import sys + import lab as B import tensorflow as tf import tensorflow_probability as tfp @@ -13,7 +15,11 @@ def take_along_axis(a: _Numeric, index: _Numeric, axis: int = 0) -> _Numeric: # """ Gathers elements of `a` along `axis` at `index` locations. """ - return tf.gather(a, B.flatten(index), axis=axis) + if sys.version_info[:2] <= (3, 9): + index = tf.cast(index, tf.int32) + return tf.experimental.numpy.take_along_axis( + a, index, axis=axis + ) # the absence of explicit cast to int64 causes an error for Python 3.9 and below @dispatch @@ -164,10 +170,7 @@ def complex_like(reference: B.TFNumeric): """ Return `complex` dtype of a backend based on the reference. """ - if B.dtype(reference) == tf.float32: - return tf.complex64 - else: - return tf.complex128 + return B.promote_dtypes(tf.complex64, reference.dtype) @dispatch @@ -251,3 +254,16 @@ def dtype_bool(reference: B.TFRandomState): # type: ignore Return `bool` dtype of a backend based on the reference. """ return tf.bool + + +@dispatch +def bool_like(reference: B.NPNumeric): + """ + Return the type of the reference if it is of boolean type. + Otherwise return `bool` dtype of a backend based on the reference. + """ + reference_dtype = reference.dtype + if reference_dtype.is_bool: + return reference_dtype + else: + return tf.bool diff --git a/geometric_kernels/lab_extras/torch/extras.py b/geometric_kernels/lab_extras/torch/extras.py index 15860ec5..b16a5759 100644 --- a/geometric_kernels/lab_extras/torch/extras.py +++ b/geometric_kernels/lab_extras/torch/extras.py @@ -14,7 +14,7 @@ def take_along_axis(a: Union[_Numeric, B.Numeric], index: _Numeric, axis: int = """ if not torch.is_tensor(a): a = torch.tensor(a).to(index.device) # type: ignore - return torch.index_select(a, axis, B.flatten(index)) + return torch.take_along_dim(a, index.long(), axis) # long is required by torch @dispatch @@ -171,10 +171,7 @@ def complex_like(reference: B.TorchNumeric): """ Return `complex` dtype of a backend based on the reference. """ - if B.dtype(reference) == torch.float: - return torch.cfloat - else: - return torch.cdouble + return B.promote_dtypes(torch.cfloat, reference.dtype) @dispatch @@ -258,3 +255,16 @@ def dtype_bool(reference: B.TorchRandomState): # type: ignore Return `bool` dtype of a backend based on the reference. """ return torch.bool + + +@dispatch +def bool_like(reference: B.NPNumeric): + """ + Return the type of the reference if it is of boolean type. + Otherwise return `bool` dtype of a backend based on the reference. + """ + reference_dtype = reference.dtype + if reference_dtype is torch.bool: + return reference_dtype + else: + return torch.bool diff --git a/geometric_kernels/spaces/base.py b/geometric_kernels/spaces/base.py index 6d192bf2..9d3b7ca4 100644 --- a/geometric_kernels/spaces/base.py +++ b/geometric_kernels/spaces/base.py @@ -35,9 +35,21 @@ def element_shape(self) -> List[int]: Shape of an element. Examples: - * hypersphere: [D + 1, ] - * mesh: [1, ] - * matrix Lie group: [n, n] + * :class:`~.spaces.Hypersphere`: [D + 1, ] + * :class:`~.spaces.Mesh`: [1, ] + * :class:`~.spaces.CompactMatrixLieGroup`: [n, n] + """ + raise NotImplementedError + + @abc.abstractproperty + def element_dtype(self) -> B.DType: + """ + Abstract DType of an element. + + Examples: + * :class:`~.spaces.Hypersphere`: B.Float + * :class:`~.spaces.Mesh`: B.Int + * :class:`~.spaces.SpecialUnitary`: B.Complex """ raise NotImplementedError @@ -170,7 +182,7 @@ class NoncompactSymmetricSpace(Space): Mathematically, any non-compact symmetric space can be represented as a quotient $G/H$ of a Lie group of symmetries $G$ and its compact isotropy subgroup $H$. We sometimes refer to these $G$ and $H$ in - the documentation. See mathematical details in :cite:t:`azangulov2023`. + the documentation. See mathematical details in :cite:t:`azangulov2024b`. """ @abc.abstractproperty diff --git a/geometric_kernels/spaces/circle.py b/geometric_kernels/spaces/circle.py index 4edee885..d80d4d51 100644 --- a/geometric_kernels/spaces/circle.py +++ b/geometric_kernels/spaces/circle.py @@ -143,6 +143,9 @@ class Circle(DiscreteSpectrumSpace): citing :cite:t:`borovitskiy2020`. """ + def __str__(self): + return "Circle()" + @property def dimension(self) -> int: """ @@ -185,3 +188,11 @@ def element_shape(self): [1]. """ return [1] + + @property + def element_dtype(self): + """ + :return: + B.Float. + """ + return B.Float diff --git a/geometric_kernels/spaces/eigenfunctions.py b/geometric_kernels/spaces/eigenfunctions.py index 8bdb9055..5ad7b8f6 100644 --- a/geometric_kernels/spaces/eigenfunctions.py +++ b/geometric_kernels/spaces/eigenfunctions.py @@ -102,6 +102,7 @@ def weighted_outerproduct( if is_complex(sum_phi_phi_for_level): sum_phi_phi_for_level = B.cast(complex_like(weights), sum_phi_phi_for_level) + weights = B.cast(complex_like(weights), weights) else: sum_phi_phi_for_level = B.cast(B.dtype(weights), sum_phi_phi_for_level) @@ -128,6 +129,7 @@ def weighted_outerproduct_diag( if is_complex(phi_product_diag): phi_product_diag = B.cast(complex_like(weights), phi_product_diag) + weights = B.cast(complex_like(weights), weights) else: phi_product_diag = B.cast(B.dtype(weights), phi_product_diag) diff --git a/geometric_kernels/spaces/graph.py b/geometric_kernels/spaces/graph.py index 0ac1ad19..bfdf37f5 100644 --- a/geometric_kernels/spaces/graph.py +++ b/geometric_kernels/spaces/graph.py @@ -60,6 +60,10 @@ def __init__(self, adjacency_matrix: B.Numeric, normalize_laplacian: bool = Fals self.cache: Dict[int, Tuple[B.Numeric, B.Numeric]] = {} self._checks(adjacency_matrix) self._set_laplacian(adjacency_matrix, normalize_laplacian) # type: ignore + self._normalized = normalize_laplacian + + def __str__(self): + return f"Graph({self.num_vertices}, {'normalized' if self._normalized else 'unnormalized'})" @staticmethod def _checks(adjacency): @@ -67,11 +71,10 @@ def _checks(adjacency): Checks if `adjacency` is a square symmetric matrix. """ assert ( - len(B.shape(adjacency)) == 2 and adjacency.shape[0] == adjacency.shape[1] + len(adjacency.shape) == 2 and adjacency.shape[0] == adjacency.shape[1] ), "Matrix is not square." - # this is more efficient than (adj == adj.T).all() - assert not B.any(adjacency != B.T(adjacency)), "Adjacency is not symmetric." + assert not B.any(adjacency != B.T(adjacency)), "Adjacency is not symmetric" @property def dimension(self) -> int: @@ -115,6 +118,9 @@ def get_eigensystem(self, num): :return: A tuple of eigenvectors [n, num], eigenvalues [num, 1]. """ + assert ( + num <= self.num_vertices + ), "Number of eigenpairs cannot exceed the number of vertices" if num not in self.cache: evals, evecs = eigenpairs(self._laplacian, num) @@ -184,3 +190,11 @@ def element_shape(self): [1]. """ return [1] + + @property + def element_dtype(self): + """ + :return: + B.Int. + """ + return B.Int diff --git a/geometric_kernels/spaces/hyperbolic.py b/geometric_kernels/spaces/hyperbolic.py index 61b57ebc..e0417f92 100644 --- a/geometric_kernels/spaces/hyperbolic.py +++ b/geometric_kernels/spaces/hyperbolic.py @@ -6,13 +6,12 @@ import lab as B from beartype.typing import Optional -from geometric_kernels.lab_extras import ( - complex_like, - create_complex, - dtype_double, - from_numpy, -) +from geometric_kernels.lab_extras import complex_like, create_complex, dtype_double from geometric_kernels.spaces.base import NoncompactSymmetricSpace +from geometric_kernels.utils.manifold_utils import ( + hyperbolic_distance, + minkowski_inner_product, +) class Hyperbolic(NoncompactSymmetricSpace, gs.geometry.hyperboloid.Hyperboloid): @@ -42,17 +41,20 @@ class Hyperbolic(NoncompactSymmetricSpace, gs.geometry.hyperboloid.Hyperboloid): is a quotient G/H. For the hyperbolic space $\mathbb{H}_n$, the group of symmetries $G$ is the proper Lorentz group $SO(1, n)$, while the isotropy subgroup $H$ is the special orthogonal group $SO(n)$. See the - mathematical details in :cite:t:`azangulov2023`. + mathematical details in :cite:t:`azangulov2024b`. .. admonition:: Citation If you use this GeometricKernels space in your research, please consider - citing :cite:t:`azangulov2023`. + citing :cite:t:`azangulov2024b`. """ def __init__(self, dim=2): super().__init__(dim=dim) + def __str__(self): + return f"Hyperbolic({self.dimension})" + @property def dimension(self) -> int: """ @@ -64,76 +66,15 @@ def distance( self, x1: B.Numeric, x2: B.Numeric, diag: Optional[bool] = False ) -> B.Numeric: """ - Compute the hyperbolic distance between `x1` and `x2`. - - The code is a reimplementation of - `geomstats.geometry.hyperboloid.HyperbolicMetric` for `lab`. - - :param x1: - An [N, n+1]-shaped array of points in the hyperbolic space. - :param x2: - An [M, n+1]-shaped array of points in the hyperbolic space. - :param diag: - If True, compute elementwise distance. Requires N = M. - - Default False. - - :return: - An [N, M]-shaped array if diag=False or [N,]-shaped array - if diag=True. + Calls :func:`~.hyperbolic_distance` on the same inputs. """ - if diag: - # Compute a pointwise distance between `x1` and `x2` - x1_ = x1 - x2_ = x2 - else: - if B.rank(x1) == 1: - x1 = B.expand_dims(x1) - if B.rank(x2) == 1: - x2 = B.expand_dims(x2) - - # compute pairwise distance between arrays of points `x1` and `x2` - # `x1` (N, n+1) - # `x2` (M, n+1) - x1_ = B.tile(x1[..., None, :], 1, x2.shape[0], 1) # (N, M, n+1) - x2_ = B.tile(x2[None], x1.shape[0], 1, 1) # (N, M, n+1) - - sq_norm_1 = self.inner_product(x1_, x1_) - sq_norm_2 = self.inner_product(x2_, x2_) - inner_prod = self.inner_product(x1_, x2_) - - cosh_angle = -inner_prod / B.sqrt(sq_norm_1 * sq_norm_2) - - one = B.cast(B.dtype(cosh_angle), from_numpy(cosh_angle, [1.0])) - large_constant = B.cast(B.dtype(cosh_angle), from_numpy(cosh_angle, [1e24])) - - # clip values into [1.0, 1e24] - cosh_angle = B.where(cosh_angle < one, one, cosh_angle) - cosh_angle = B.where(cosh_angle > large_constant, large_constant, cosh_angle) - - dist = B.log(cosh_angle + B.sqrt(cosh_angle**2 - 1)) # arccosh - dist = B.cast(B.dtype(x1_), dist) - return dist + return hyperbolic_distance(x1, x2, diag) def inner_product(self, vector_a, vector_b): r""" - Computes the Minkowski inner product of vectors. - - .. math:: \langle a, b \rangle = a_0 b_0 - a_1 b_1 - \ldots - a_n b_n. - - :param vector_a: - An [..., n+1]-shaped array of points in the hyperbolic space. - :param vector_b: - An [..., n+1]-shaped array of points in the hyperbolic space. - - :return: - An [...,]-shaped array of inner products. + Calls :func:`~.minkowski_inner_product` on `vector_a` and `vector_b`. """ - q = self.dimension - p = 1 - diagonal = from_numpy(vector_a, [-1.0] * p + [1.0] * q) # (n+1) - diagonal = B.cast(B.dtype(vector_a), diagonal) - return B.einsum("...i,...i->...", diagonal * vector_a, vector_b) + return minkowski_inner_product(vector_a, vector_b) def inv_harish_chandra(self, lam: B.Numeric) -> B.Numeric: lam = B.squeeze(lam, -1) @@ -205,7 +146,7 @@ def random_phases(self, key, num): def random(self, key, number): """ - Geomstats-based non-uniform random sampling. + Non-uniform random sampling, reimplements the algorithm from geomstats. Always returns [N, n+1] float64 array of the `key`'s backend. @@ -216,10 +157,15 @@ def random(self, key, number): Number of samples to draw. :return: - An array of `number` uniformly random samples on the space. + An array of `number` random samples on the space. """ - return key, B.cast(dtype_double(key), self.random_point(number)) + key, samples = B.rand(key, dtype_double(key), number, self.dim) + + samples = 2.0 * (samples - 0.5) + + coord_0 = B.sqrt(1.0 + B.sum(samples**2, axis=-1)) + return key, B.concat(coord_0[..., None], samples, axis=-1) def element_shape(self): """ @@ -227,3 +173,11 @@ def element_shape(self): [n+1]. """ return [self.dimension + 1] + + @property + def element_dtype(self): + """ + :return: + B.Float. + """ + return B.Float diff --git a/geometric_kernels/spaces/hypercube_graph.py b/geometric_kernels/spaces/hypercube_graph.py index 0f58b7a6..2034fead 100644 --- a/geometric_kernels/spaces/hypercube_graph.py +++ b/geometric_kernels/spaces/hypercube_graph.py @@ -195,6 +195,9 @@ def __init__(self, dim: int): raise ValueError("dim must be a positive integer.") self.dim = dim + def __str__(self): + return f"HypercubeGraph({self.dim})" + @property def dimension(self) -> int: """ @@ -269,3 +272,11 @@ def element_shape(self): [d]. """ return [self.dimension] + + @property + def element_dtype(self): + """ + :return: + B.Bool. + """ + return B.Bool diff --git a/geometric_kernels/spaces/hypersphere.py b/geometric_kernels/spaces/hypersphere.py index f1674102..5aa1870a 100644 --- a/geometric_kernels/spaces/hypersphere.py +++ b/geometric_kernels/spaces/hypersphere.py @@ -106,10 +106,7 @@ def _addition_theorem_diag(self, X: B.Numeric, **kwargs) -> B.Numeric: @property def num_eigenfunctions(self) -> int: if self._num_eigenfunctions is None: - J = 0 - for level in range(self.num_levels): - J += num_harmonics(self.dim + 1, level) - self._num_eigenfunctions = J + self._num_eigenfunctions = sum(self.num_eigenfunctions_per_level) return self._num_eigenfunctions @property @@ -154,6 +151,9 @@ def __init__(self, dim: int): super().__init__(dim=dim) self.dim = dim + def __str__(self): + return f"Hypersphere({self.dim})" + @property def dimension(self) -> int: """ @@ -257,3 +257,11 @@ def element_shape(self): [d+1]. """ return [self.dimension + 1] + + @property + def element_dtype(self): + """ + :return: + B.Float. + """ + return B.Float diff --git a/geometric_kernels/spaces/lie_groups.py b/geometric_kernels/spaces/lie_groups.py index 9a79efa4..0ee11441 100644 --- a/geometric_kernels/spaces/lie_groups.py +++ b/geometric_kernels/spaces/lie_groups.py @@ -43,7 +43,7 @@ class WeylAdditionTheorem(EigenfunctionsWithAdditionTheorem): is representation-theoretic: they are proportional to *characters* of irreducible unitary representations of the group. These characters, in their turn, can be algebraically computed using the *Weyl character formula*. See - :cite:t:`azangulov2022` for the mathematical details behind this class. + :cite:t:`azangulov2024a` for the mathematical details behind this class. :param n: The order of the Lie group, e.g. for SO(5) this is 5, for SU(3) this is 3. @@ -92,6 +92,7 @@ def __init__(self, n: int, num_levels: int, compute_characters: bool = True): self._characters = [ self._compute_character(n, signature) for signature in self._signatures ] + self._num_eigenfunctions: Optional[int] = None # To be computed when needed. @abc.abstractmethod def _generate_signatures(self, num_levels: int) -> List[Tuple[int, ...]]: @@ -195,15 +196,15 @@ def _difference(self, X: B.Numeric, X2: B.Numeric) -> B.Numeric: .. note:: Doing X1[j, :, :] * inv(X2[i, :, :]) is as permissible as doing inv(X2[i, :, :]) * X1[j, :, :] which is actually used in - :cite:t:`azangulov2022`. This is because $\chi(x y x^{-1})=\chi(y)$ + :cite:t:`azangulov2024a`. This is because $\chi(x y x^{-1})=\chi(y)$ which implies that $\chi(x y) = \chi(y x)$. """ X2_inv = self.inverse(X2) X_ = B.tile(X[..., None, :, :], 1, X2_inv.shape[0], 1, 1) # (N, N2, n, n) X2_inv_ = B.tile(X2_inv[None, ..., :, :], X.shape[0], 1, 1, 1) # (N, N2, n, n) - diff = B.matmul(X_, X2_inv_).reshape( - X.shape[0], X2_inv.shape[0], X.shape[-1], X.shape[-1] + diff = B.reshape( + B.matmul(X_, X2_inv_), X.shape[0], X2_inv.shape[0], X.shape[-1], X.shape[-1] ) # (N, N2, n, n) return diff @@ -216,7 +217,7 @@ def _addition_theorem( Laplace-Beltrami eigenfunctions that correspond to this level (representation). Uses the fact that such sums are equal to the character of the representation multiplied by the dimension of that - representation. See :cite:t:`azangulov2022` for mathematical details. + representation. See :cite:t:`azangulov2024a` for mathematical details. :param X: An [N, n, n]-shaped array, a batch of N matrices of size nxn. @@ -266,7 +267,9 @@ def num_levels(self) -> int: @property def num_eigenfunctions(self) -> int: - return B.sum(self.num_eigenfunctions_per_level) + if self._num_eigenfunctions is None: + self._num_eigenfunctions = sum(self.num_eigenfunctions_per_level) + return self._num_eigenfunctions @property def num_eigenfunctions_per_level(self) -> List[int]: diff --git a/geometric_kernels/spaces/mesh.py b/geometric_kernels/spaces/mesh.py index 7b744ab6..122806f7 100644 --- a/geometric_kernels/spaces/mesh.py +++ b/geometric_kernels/spaces/mesh.py @@ -66,6 +66,9 @@ def __init__(self, vertices: np.ndarray, faces: np.ndarray): self._eigenfunctions = None self.cache: Dict[int, Tuple[np.ndarray, np.ndarray]] = {} + def __str__(self): + return f"Mesh({self.num_vertices})" + def get_eigensystem(self, num: int) -> Tuple[np.ndarray, np.ndarray]: """ Returns the first `num` eigenvalues and eigenvectors of the `robust @@ -214,3 +217,11 @@ def element_shape(self): [1]. """ return [1] + + @property + def element_dtype(self): + """ + :return: + B.Int. + """ + return B.Int diff --git a/geometric_kernels/spaces/product.py b/geometric_kernels/spaces/product.py index f278be7a..c2b7de1c 100644 --- a/geometric_kernels/spaces/product.py +++ b/geometric_kernels/spaces/product.py @@ -15,7 +15,7 @@ import numpy as np from beartype.typing import List, Optional, Tuple -from geometric_kernels.lab_extras import from_numpy, int_like +from geometric_kernels.lab_extras import from_numpy, int_like, take_along_axis from geometric_kernels.spaces.base import DiscreteSpectrumSpace from geometric_kernels.spaces.eigenfunctions import Eigenfunctions from geometric_kernels.utils.product import make_product, project_product @@ -154,7 +154,11 @@ class ProductEigenfunctions(Eigenfunctions): Levels correspond to tuples of levels of the factors. :param element_shapes: - Shapes of the elements in each of the factor spaces. + Shapes of the elements in each factor. Can be obtained as properties + `space.element_shape` of any given factor `space`. + :param element_dtypes: + Abstract lab data types of the elements in each factor. Can be obtained + as properties `space.element_dtype` of any given factor `space`. :param eigenindicies: A [L, S]-shaped array, where `S` is the number of factor spaces and `L` is the number of levels, such that `eigenindicies[i, :]` are the indices @@ -184,12 +188,14 @@ class ProductEigenfunctions(Eigenfunctions): def __init__( self, element_shapes: List[List[int]], + element_dtypes: List[B.DType], eigenindicies: B.Numeric, *eigenfunctions: Eigenfunctions, dimension_indices: B.Numeric = None, ): self._num_levels = eigenindicies.shape[0] self.element_shapes = element_shapes + self.element_dtypes = element_dtypes dimensions = [math.prod(element_shape) for element_shape in self.element_shapes] if dimension_indices is None: self.dimension_indices = [] @@ -240,7 +246,9 @@ def __call__(self, X: B.Numeric, **kwargs) -> B.Numeric: :return: An [N, J]-shaped array, where `J` is the number of eigenfunctions. """ - Xs = project_product(X, self.dimension_indices, self.element_shapes) + Xs = project_product( + X, self.dimension_indices, self.element_shapes, self.element_dtypes + ) # List of S arrays, each of shape [N, *element_shape_s] factor_eigenfunction_values = [ eigenfunction(X, **kwargs) # [N, Js], Js different for each factor @@ -256,13 +264,13 @@ def __call__(self, X: B.Numeric, **kwargs) -> B.Numeric: *( factor_eigenfunction_values[s][ :, self._eigenfunctionindices[j, s] - ] + ] # [N,] for s in range(len(m_idx)) ), axis=-1, - ), + ), # [N, S] axis=1, - ) + ) # [N,] ) eigenfunction_values = B.stack(*eigenfunction_values, axis=-1) # [N, J] @@ -281,46 +289,51 @@ def phi_product( ) -> B.Numeric: if X2 is None: X2 = X - Xs = project_product(X, self.dimension_indices, self.element_shapes) - Xs2 = project_product(X2, self.dimension_indices, self.element_shapes) + Xs = project_product( + X, self.dimension_indices, self.element_shapes, self.element_dtypes + ) # List of S arrays, each of shape [N, *element_shape_s] + Xs2 = project_product( + X2, self.dimension_indices, self.element_shapes, self.element_dtypes + ) # List of S arrays, each of shape [N2, *element_shape_s] + + factor_phi_products = [ + take_along_axis( + eigenfunction.phi_product(X1, X2, **kwargs), # [N, N2, L_s] + from_numpy(X1, self.eigenindicies[None, None, :, s]), + -1, + ) # [N, N2, L] + for s, (eigenfunction, X1, X2) in enumerate( + zip(self.eigenfunctions, Xs, Xs2) + ) + ] + common_dtype = B.promote_dtypes(*[B.dtype(x) for x in factor_phi_products]) phis = B.stack( - *[ - eigenfunction.phi_product(X1, X2, **kwargs) - for eigenfunction, X1, X2 in zip(self.eigenfunctions, Xs, Xs2) - ], - axis=-1, - ) # [N, N2, LFactor, S] where `LFactor` is self.eigenfunctions[0].num_levels + *[B.cast(common_dtype, x) for x in factor_phi_products], axis=-1 + ) # [N, N2, L, S] - prod_phis = phis[ - :, - :, - self.eigenindicies, - B.range(self.eigenindicies.shape[1]), - ].prod( - axis=-1 - ) # [N, N2, LFactor, S] -> [N, N2, L] + prod_phis = B.prod(phis, axis=-1) # [N, N2, L, S] -> [N, N2, L] return prod_phis def phi_product_diag(self, X: B.Numeric, **kwargs): - Xs = project_product(X, self.dimension_indices, self.element_shapes) + Xs = project_product( + X, self.dimension_indices, self.element_shapes, self.element_dtypes + ) phis = B.stack( *[ - eigenfunction.phi_product_diag(X1, **kwargs) - for eigenfunction, X1 in zip(self.eigenfunctions, Xs) + take_along_axis( + eigenfunction.phi_product_diag(X1, **kwargs), # [N, L_s] + from_numpy(X1, self.eigenindicies[None, :, s]), + -1, + ) # [N, L] + for s, (eigenfunction, X1) in enumerate(zip(self.eigenfunctions, Xs)) ], axis=-1, - ) # [N, LFactor, S] where `LFactor` is self.eigenfunctions[0].num_levels + ) # [N, L, S] - prod_phis = phis[ - :, - self.eigenindicies, - B.range(self.eigenindicies.shape[1]), - ].prod( - axis=-1 - ) # [N, LFactor, S] -> [N, L] + prod_phis = B.prod(phis, axis=-1) # [N, L] return prod_phis @@ -440,6 +453,9 @@ def __init__( self.factor_space_eigenindices, ].sum(axis=1) + def __str__(self): + return f"ProductDiscreteSpectrumSpace({', '.join(str(space) for space in self.factor_spaces)})" + @property def dimension(self) -> int: """ @@ -487,6 +503,7 @@ def get_eigenfunctions(self, num: int) -> Eigenfunctions: return ProductEigenfunctions( [space.element_shape for space in self.factor_spaces], + [space.element_dtype for space in self.factor_spaces], self.factor_space_eigenindices[:num], *factor_space_eigenfunctions, ) @@ -534,3 +551,11 @@ def element_shape(self): Sum of the products of the element shapes of the factor spaces. """ return sum(math.prod(space.element_shape) for space in self.factor_spaces) + + @property + def element_dtype(self): + """ + :return: + B.Numeric. + """ + return B.Numeric diff --git a/geometric_kernels/spaces/so.py b/geometric_kernels/spaces/so.py index 0ed3cfa7..6ea1969a 100644 --- a/geometric_kernels/spaces/so.py +++ b/geometric_kernels/spaces/so.py @@ -14,7 +14,15 @@ import numpy as np from beartype.typing import List, Tuple -from geometric_kernels.lab_extras import dtype_double, from_numpy, qr, take_along_axis +from geometric_kernels.lab_extras import ( + complex_conj, + complex_like, + create_complex, + dtype_double, + from_numpy, + qr, + take_along_axis, +) from geometric_kernels.spaces.eigenfunctions import Eigenfunctions from geometric_kernels.spaces.lie_groups import ( CompactMatrixLieGroup, @@ -129,7 +137,7 @@ def _torus_representative(self, X: B.Numeric) -> B.Numeric: real = (trace - 1) / 2 zeros = real * 0 imag = B.sqrt(B.maximum(1 - real * real, zeros)) - gamma = real + 1j * imag + gamma = create_complex(real, imag) elif self.n % 2 == 1: # In SO(2n+1) the torus representative is determined by the (unordered) non-trivial eigenvalues eigvals = B.eig(X, False) @@ -147,17 +155,30 @@ def _torus_representative(self, X: B.Numeric) -> B.Numeric: -1, ) # c is a matrix transforming x into its canonical form (with 2x2 blocks) - c = 0 * eigvecs - c[..., ::2] = eigvecs[..., ::2] + eigvecs[..., 1::2] - c[..., 1::2] = eigvecs[..., ::2] - eigvecs[..., 1::2] + c = B.reshape( + B.stack( + eigvecs[..., ::2] + eigvecs[..., 1::2], + eigvecs[..., ::2] - eigvecs[..., 1::2], + axis=-1, + ), + *eigvecs.shape[:-1], + -1, + ) # eigenvectors calculated by LAPACK are either real or purely imaginary, make everything real # WARNING: might depend on the implementation of the eigendecomposition! - c = c.real + c.imag + c = B.real(c) + B.imag(c) # normalize s.t. det(c)≈±1, probably unnecessary c /= math.sqrt(2) - eigvals[..., 0] = B.power(eigvals[..., 0], B.sign(B.det(c))) + eigvals = B.concat( + B.expand_dims( + B.power(eigvals[..., 0], B.cast(complex_like(c), B.sign(B.det(c)))), + axis=-1, + ), + eigvals[..., 1:], + axis=-1, + ) gamma = eigvals[..., ::2] - gamma = B.concat(gamma, gamma.conj(), axis=-1) + gamma = B.concat(gamma, complex_conj(gamma), axis=-1) return gamma def inverse(self, X: B.Numeric) -> B.Numeric: @@ -172,7 +193,7 @@ class SOCharacter(LieGroupCharacter): These are polynomials whose coefficients are precomputed and stored in a file. By default, there are 20 precomputed characters for n from 3 to 8. - If you want more, use the `utils/compute_characters.py` script. + If you want more, use the `compute_characters.py` script. :param n: The order n of the SO(n) group. @@ -207,7 +228,9 @@ def _load(self): def __call__(self, gammas: B.Numeric) -> B.Numeric: char_val = B.zeros(B.dtype(gammas), *gammas.shape[:-1]) for coeff, monom in zip(self.coeffs, self.monoms): - char_val += coeff * B.prod(gammas ** from_numpy(gammas, monom), axis=-1) + char_val += coeff * B.prod( + gammas ** B.cast(B.dtype(gammas), from_numpy(gammas, monom)), axis=-1 + ) return char_val @@ -238,7 +261,7 @@ class SpecialOrthogonal(CompactMatrixLieGroup): .. admonition:: Citation If you use this GeometricKernels space in your research, please consider - citing :cite:t:`azangulov2022`. + citing :cite:t:`azangulov2024a`. """ def __init__(self, n: int): @@ -249,6 +272,9 @@ def __init__(self, n: int): self.rank = n // 2 super().__init__() + def __str__(self): + return f"SpecialOrthogonal({self.n})" + @property def dimension(self) -> int: """ @@ -323,8 +349,9 @@ def random(self, key: B.RandomState, number: int): r_diag_sign = B.sign(B.einsum("...ii->...i", r)) q *= r_diag_sign[:, None] q_det_sign = B.sign(B.det(q)) - q[:, :, 0] *= q_det_sign[:, None] - return key, q + q_new = q[:, :, 0] * q_det_sign[:, None] + q_new = B.concat(q_new[:, :, None], q[:, :, 1:], axis=-1) + return key, q_new @property def element_shape(self): @@ -333,3 +360,11 @@ def element_shape(self): [n, n]. """ return [self.n, self.n] + + @property + def element_dtype(self): + """ + :return: + B.Float. + """ + return B.Float diff --git a/geometric_kernels/spaces/spd.py b/geometric_kernels/spaces/spd.py index 0c329bd2..13ec3952 100644 --- a/geometric_kernels/spaces/spd.py +++ b/geometric_kernels/spaces/spd.py @@ -43,17 +43,20 @@ class SymmetricPositiveDefiniteMatrices( matrices $SPD(n)$, the group of symmetries $G$ is the identity component $GL(n)_+$ of the general linear group $GL(n)$, while the isotropy subgroup $H$ is the special orthogonal group $SO(n)$. See the - mathematical details in :cite:t:`azangulov2023`. + mathematical details in :cite:t:`azangulov2024b`. .. admonition:: Citation If you use this GeometricKernels space in your research, please consider - citing :cite:t:`azangulov2023`. + citing :cite:t:`azangulov2024b`. """ def __init__(self, n): super().__init__(n) + def __str__(self): + return f"SymmetricPositiveDefiniteMatrices({self.n})" + @property def dimension(self) -> int: """ @@ -117,7 +120,7 @@ def power_function(self, lam, g, h): def random(self, key, number): """ - Geomstats-based non-uniform random sampling. + Non-uniform random sampling, reimplements the algorithm from geomstats. Always returns [N, n, n] float64 array of the `key`'s backend. @@ -128,9 +131,14 @@ def random(self, key, number): Number of samples to draw. :return: - An array of `number` uniformly random samples on the space. + An array of `number` random samples on the space. """ - return key, B.cast(dtype_double(key), self.random_point(number)) + + key, mat = B.rand(key, dtype_double(key), number, self.n, self.n) + mat = 2 * mat - 1 + mat_symm = 0.5 * (mat + B.transpose(mat, (0, 2, 1))) + + return key, B.expm(mat_symm) @property def element_shape(self): @@ -139,3 +147,11 @@ def element_shape(self): [n, n]. """ return [self.n, self.n] + + @property + def element_dtype(self): + """ + :return: + B.Float. + """ + return B.Float diff --git a/geometric_kernels/spaces/su.py b/geometric_kernels/spaces/su.py index 24c872f8..5fc405df 100644 --- a/geometric_kernels/spaces/su.py +++ b/geometric_kernels/spaces/su.py @@ -15,6 +15,7 @@ from geometric_kernels.lab_extras import ( complex_conj, + complex_like, create_complex, dtype_double, from_numpy, @@ -107,7 +108,7 @@ class SUCharacter(LieGroupCharacter): These are polynomials whose coefficients are precomputed and stored in a file. By default, there are 20 precomputed characters for n from 2 to 6. - If you want more, use the `utils/compute_characters.py` script. + If you want more, use the `compute_characters.py` script. :param n: The order n of the SO(n) group. @@ -142,7 +143,12 @@ def _load(self): def __call__(self, gammas: B.Numeric) -> B.Numeric: char_val = B.zeros(B.dtype(gammas), *gammas.shape[:-1]) for coeff, monom in zip(self.coeffs, self.monoms): - char_val += coeff * B.prod(gammas ** from_numpy(gammas, monom), axis=-1) + char_val += coeff * B.prod( + B.power( + gammas, B.cast(complex_like(gammas), from_numpy(gammas, monom)) + ), + axis=-1, + ) return char_val @@ -171,7 +177,7 @@ class SpecialUnitary(CompactMatrixLieGroup): .. admonition:: Citation If you use this GeometricKernels space in your research, please consider - citing :cite:t:`azangulov2022`. + citing :cite:t:`azangulov2024a`. """ def __init__(self, n: int): @@ -182,6 +188,9 @@ def __init__(self, n: int): self.rank = n - 1 super().__init__() + def __str__(self): + return f"SpecialUnitary({self.n})" + @property def dimension(self) -> int: """ @@ -242,12 +251,17 @@ def random(self, key: B.RandomState, number: int): h = create_complex(real, imag) / B.sqrt(2) q, r = qr(h, mode="complete") r_diag = B.einsum("...ii->...i", r) - r_diag_inv_phase = complex_conj(r_diag / B.abs(r_diag)) + r_diag_inv_phase = complex_conj( + r_diag / B.cast(B.dtype(r_diag), B.abs(r_diag)) + ) q *= r_diag_inv_phase[:, None] q_det = B.det(q) - q_det_inv_phase = complex_conj((q_det / B.abs(q_det))) - q[:, :, 0] *= q_det_inv_phase[:, None] - return key, q + q_det_inv_phase = complex_conj( + (q_det / B.cast(B.dtype(q_det), B.abs(q_det))) + ) + q_new = q[:, :, 0] * q_det_inv_phase[:, None] + q_new = B.concat(q_new[:, :, None], q[:, :, 1:], axis=-1) + return key, q_new @property def element_shape(self): @@ -256,3 +270,11 @@ def element_shape(self): [n, n]. """ return [self.n, self.n] + + @property + def element_dtype(self): + """ + :return: + B.Complex. + """ + return B.Complex diff --git a/geometric_kernels/utils/kernel_formulas/__init__.py b/geometric_kernels/utils/kernel_formulas/__init__.py new file mode 100644 index 00000000..82bca600 --- /dev/null +++ b/geometric_kernels/utils/kernel_formulas/__init__.py @@ -0,0 +1,14 @@ +from geometric_kernels.utils.kernel_formulas.euclidean import ( + euclidean_matern_12_kernel, + euclidean_matern_32_kernel, + euclidean_matern_52_kernel, + euclidean_rbf_kernel, +) +from geometric_kernels.utils.kernel_formulas.hyperbolic import ( + hyperbolic_heat_kernel_even, + hyperbolic_heat_kernel_odd, +) +from geometric_kernels.utils.kernel_formulas.hypercube_graph import ( + hypercube_graph_heat_kernel, +) +from geometric_kernels.utils.kernel_formulas.spd import spd_heat_kernel_2x2 diff --git a/geometric_kernels/utils/kernel_formulas/euclidean.py b/geometric_kernels/utils/kernel_formulas/euclidean.py new file mode 100644 index 00000000..f76f347d --- /dev/null +++ b/geometric_kernels/utils/kernel_formulas/euclidean.py @@ -0,0 +1,103 @@ +""" +Implements the standard formulas for the RBF kernel and some Matérn kernels. + +The implementation is provided mainly for testing purposes. +""" + +from math import sqrt + +import lab as B +from beartype.typing import Optional + + +def euclidean_matern_12_kernel( + r: B.Numeric, + lengthscale: Optional[float] = 1.0, +): + """ + Analytic formula for the Matérn 1/2 kernel on R^d, as a function of + distance `r` between inputs. + + :param r: + A batch of distances, an array of shape [...]. + :param lengthscale: + The length scale of the kernel, defaults to 1. + + :return: + The kernel values evaluated at `r`, an array of shape [...]. + """ + + assert B.all(r >= 0.0) + + return B.exp(-r / lengthscale) + + +def euclidean_matern_32_kernel( + r: B.Numeric, + lengthscale: Optional[float] = 1.0, +): + """ + Analytic formula for the Matérn 3/2 kernel on R^d, as a function of + distance `r` between inputs. + + :param r: + A batch of distances, an array of shape [...]. + :param lengthscale: + The length scale of the kernel, defaults to 1. + + :return: + The kernel values evaluated at `r`, an array of shape [...]. + """ + + assert B.all(r >= 0.0) + + sqrt3 = sqrt(3.0) + r = r / lengthscale + return (1.0 + sqrt3 * r) * B.exp(-sqrt3 * r) + + +def euclidean_matern_52_kernel( + r: B.Numeric, + lengthscale: Optional[float] = 1.0, +): + """ + Analytic formula for the Matérn 5/2 kernel on R^d, as a function of + distance `r` between inputs. + + :param r: + A batch of distances, an array of shape [...]. + :param lengthscale: + The length scale of the kernel, defaults to 1. + + :return: + The kernel values evaluated at `r`, an array of shape [...]. + """ + + assert B.all(r >= 0.0) + + sqrt5 = sqrt(5.0) + r = r / lengthscale + return (1.0 + sqrt5 * r + 5.0 / 3.0 * (r**2)) * B.exp(-sqrt5 * r) + + +def euclidean_rbf_kernel( + r: B.Numeric, + lengthscale: Optional[float] = 1.0, +): + """ + Analytic formula for the RBF kernel on R^d, as a function of + distance `r` between inputs. + + :param r: + A batch of distances, an array of shape [...]. + :param lengthscale: + The length scale of the kernel, defaults to 1. + + :return: + The kernel values evaluated at `r`, an array of shape [...]. + """ + + assert B.all(r >= 0.0) + + r = r / lengthscale + return B.exp(-0.5 * r**2) diff --git a/geometric_kernels/utils/kernel_formulas/hyperbolic.py b/geometric_kernels/utils/kernel_formulas/hyperbolic.py new file mode 100644 index 00000000..575c0612 --- /dev/null +++ b/geometric_kernels/utils/kernel_formulas/hyperbolic.py @@ -0,0 +1,198 @@ +""" +Implements alternative formulas for the heat kernel on the hyperbolic manifold. + +More specifically, the function :func:`hyperbolic_heat_kernel_odd` implements +analytic formulas for the heat kernel on odd-dimensional hyperbolic spaces. The +function :func:`hyperbolic_heat_kernel_even` implements a semi-analytic formula +for the heat kernel on even-dimensional hyperbolic spaces. + +The implementation is provided mainly for testing purposes. Hypothetically, the +odd-dimensional formula could be used in practice, but the even-dimensional one +is not recommended due to its inefficiency and numerical instability. +""" + +import lab as B +import numpy as np +import scipy +from beartype.typing import Optional + +from geometric_kernels.lab_extras import cosh, sinh +from geometric_kernels.utils.manifold_utils import hyperbolic_distance + + +def hyperbolic_heat_kernel_odd( + dim: int, + t: float, + X: B.Numeric, + X2: Optional[B.Numeric] = None, +) -> B.Numeric: + """ + The analytic formula for the heat kernel on the hyperbolic space of odd + dimension, normalized to have k(x, x) = 1 for all x. + + :param t: + The time parameter, a positive float. + :param X: + A batch of inputs, an array of shape [N, dim+1]. + :param X2: + A batch of inputs, an array of shape [N2, dim+1]. If None, defaults to X. + + :return: + The kernel matrix, an array of shape [N, N2]. The kernel is normalized, + i.e. k(x, x) = 1 for all x. + """ + + if dim % 2 == 0: + raise ValueError( + "This function is only defined for odd-dimensional hyperbolic spaces. For even-dimensional spaces, use `hyperbolic_heat_kernel_even`." + ) + + if X2 is None: + X2 = X + + dists = hyperbolic_distance(X, X2) + dists_are_small = dists < 0.1 + + if dim == 3: + # This formula is taken from :cite:t:`azangulov2024b`, Equation (42). + analytic_expr = dists / sinh(dists) * B.exp(-(dists**2) / (4 * t)) + + # To get the Taylor expansion below, which gives a stable way to compute + # the kernel for small distances, use the following Mathematica code: + # > Subscript[F, 3][r_, t_] := r/Sinh[r]*Exp[-r^2/(4*t)] + # > Series[Subscript[F, 3][r, t],{r, 0, 5}] + taylor = 1 - (1 / 6 + 1 / (4 * t)) * dists**2 + + return B.where(dists_are_small, taylor, analytic_expr) + elif dim == 5: + # The following expression follows from the recursive formula in Equation + # (43) of :cite:t:`azangulov2024b`. In order to get the form below, you + # can continue the Mathematica code above with the following: + # > Subscript[U, 5][r_,t_] := -1/Sinh[r]*(D[Subscript[F, 3][rr,t ], rr] /. rr -> r) + # > Subscript[C, 5][t_] := Limit[Subscript[U, 5][r, t], r -> 0] + # > Subscript[F, 5][r_, t_] := Subscript[U, 5][r, t]/Subscript[C, 5][t] + # > Simplify[Subscript[F, 5][r, t]] + a = 3 * B.exp(-(dists**2) / (4 * t)) / (3 + 2 * t) + b = dists**2 - 2 * t + 2 * dists * t * cosh(dists) / sinh(dists) + c = 1.0 / sinh(dists) ** 2 + analytic_expr = a * b * c + + # To get the Taylor expansion below, which gives a stable way to compute + # the kernel for small distances, use the following Mathematica code: + # > Series[Subscript[F, 5][r, t], {r, 0, 5}] + taylor = 1 - (15 - 30 * t - 16 * t**2) / (20 * t * (3 + 2 * t)) * dists**2 + + return B.where(dists_are_small, taylor, analytic_expr) + elif dim == 7: + # The following expression follows from the recursive formula in Equation + # (43) of :cite:t:`azangulov2024b`. In order to get the form below, you + # can continue the Mathematica code above with the following: + # > Subscript[U, 7][r_, t_] := -1/Sinh[r]*(D[Subscript[F, 5][rr, t ], rr] /. rr -> r) + # > Subscript[C, 7][t_]:=Limit[Subscript[U, 7][r, t],r->0] + # > Subscript[F, 7][r_, t_] := Subscript[U, 7][r, t]/Subscript[C, 7][t] + # > TrigFactor[Subscript[F, 7][r,t]] + a = 15 * B.exp(-(dists**2) / (4 * t)) / (2 * (15 + 30 * t + 16 * t**2)) + b1 = -12 * t**2 * sinh(2 * dists) + b2 = (6 * t + 16 * t**2 + (8 * t**2 - 6 * t) * cosh(2 * dists)) * dists + b3 = 6 * t * sinh(2 * dists) * dists**2 + b4 = (cosh(2 * dists) - 1) * dists**3 + b = b1 + b2 + b3 + b4 + analytic_expr = a * b / sinh(dists) ** 5 + + # To get the Taylor expansion below, which gives a stable way to compute + # the kernel for small distances, use the following Mathematica code: + # > Series[Subscript[F, 7][r, t],{r, 0, 5}] + taylor = ( + 1 + - 3 + * (35 + 140 * t + 196 * t**2 + 96 * t**3) + / (28 * t * (15 + 30 * t + 16 * t**2)) + * dists**2 + ) + + return B.where(dists_are_small, taylor, analytic_expr) + else: + raise NotImplementedError( + f"Odd-dimensional hyperbolic space of dimension {dim} is not supported." + ) + + +def _hyperbolic_heat_kernel_2d_unnormalized(t: float, rho: float) -> float: + def integrand(t: float, s: float, rho: float) -> float: + result = (s + rho) * np.exp(-((s + rho) ** 2) / (4 * t)) + result /= np.sqrt((np.cosh(s + rho) - np.cosh(rho))) + return result + + integral, error = scipy.integrate.quad(lambda s: integrand(t, s, rho), 0, np.inf) + + return integral + + +def hyperbolic_heat_kernel_even( + dim: int, + t: float, + X: B.NPNumeric, + X2: Optional[B.NPNumeric] = None, +) -> B.NPNumeric: + + if dim % 2 != 0: + raise ValueError( + "This function is only defined for even-dimensional hyperbolic spaces. For odd-dimensional spaces, use `hyperbolic_heat_kernel_odd`." + ) + elif dim != 2: + # The integrand in higher dimensions may be obtained from the + # recursive formula in Equation (43) of :cite:t:`azangulov2024b`. + # + # For example, to get the integrand in dimension 4, you can use the + # following Mathematica code: + # > Subscript[F, 2][r_,t_,s_]:=((r+s)*Exp[-(r+s)^2/(4*t)])/(Cosh[r+s]-Cosh[r])^(1/2) + # > Subscript[F, 4][r_,t_,s_]:=-(1/Sinh[r])*(D[Subscript[F, 2][rr,t,s ], rr] /. rr -> r) + # > Simplify[Subscript[F, 4][r, t, s]] + # + # However, in higher dimensions, these integrands become a huge pain + # to work with. For example, the integrand in dimension 4 looks like + # + # | *** + # | * ** + # | * *** + # | * *** + # | * **** + # -|- * ------------------------- + # | * + # | * + # | * + # + # This function is very hard to numerically integrate because + # * it has a singularity at s=0 (it explodes to -inf), + # * it changes sign, + # * the domain of integration is infinite, and for large s you get 0*inf + # situations all the time (=> you have to use an asymptotic for s->inf), + # * finally, it also behaves very poorly for small rho (=> you have to + # use another asymptotic for rho->0). + # + # This is why we don't provide the integrand in higher dimensions. + # + # To whoever wants to implement it in higher dimensions in the future: + # Don't. It's not worth it. The Fourier feature approximation is already + # very accurate and, crucially, it is numerically stable. Furthermore, + # any implementation of the integrand in higher dimensions will be very + # hacky, thus diminishing its value for testing purposes. + raise NotImplementedError( + f"Even-dimensional hyperbolic space of dimension {dim} is not supported. See the comments in the code for more details on why." + ) + + if X2 is None: + X2 = X + + normalization = _hyperbolic_heat_kernel_2d_unnormalized(t, 0) + + result = np.zeros((X.shape[0], X2.shape[0])) + + # This is a very inefficient implementation, but it will do for tests. + for i, x in enumerate(X): + for j, x2 in enumerate(X2): + rho = hyperbolic_distance(x, x2).squeeze() + cur_result = _hyperbolic_heat_kernel_2d_unnormalized(t, rho) + result[i, j] = cur_result / normalization + + return result diff --git a/geometric_kernels/utils/kernel_formulas/hypercube_graph.py b/geometric_kernels/utils/kernel_formulas/hypercube_graph.py new file mode 100644 index 00000000..254d7f23 --- /dev/null +++ b/geometric_kernels/utils/kernel_formulas/hypercube_graph.py @@ -0,0 +1,52 @@ +""" +Implements the closed form expression for the heat kernel on the hypercube graph. + +The implementation is provided mainly for testing purposes. +""" + +from math import sqrt + +import lab as B +from beartype.typing import Optional + +from geometric_kernels.lab_extras import float_like +from geometric_kernels.utils.utils import hamming_distance + + +def hypercube_graph_heat_kernel( + lengthscale: B.Numeric, + X: B.Numeric, + X2: Optional[B.Numeric] = None, + normalized_laplacian: bool = True, +): + """ + Analytic formula for the heat kernel on the hypercube graph, see + Equation (14) in :cite:t:`borovitskiy2023`. + + :param lengthscale: + The length scale of the kernel, an array of shape [1]. + :param X: + A batch of inputs, an array of shape [N, d]. + :param X2: + A batch of inputs, an array of shape [N2, d]. If None, defaults to X. + + :return: + The kernel matrix, an array of shape [N, N2]. + """ + if X2 is None: + X2 = X + + assert lengthscale.shape == (1,) + assert X.ndim == 2 and X2.ndim == 2 + assert X.shape[-1] == X2.shape[-1] + + if normalized_laplacian: + d = X.shape[-1] + lengthscale = lengthscale / sqrt(d) + + # For TensorFlow, we need to explicitly cast the distances to double. + # Note: if we use B.dtype_float(X) instead of float_like(X), it gives + # float16 and TensorFlow is still complaining. + hamming_distances = B.cast(float_like(X), hamming_distance(X, X2)) + + return B.tanh(lengthscale**2 / 2) ** hamming_distances diff --git a/geometric_kernels/utils/kernel_formulas/spd.py b/geometric_kernels/utils/kernel_formulas/spd.py new file mode 100644 index 00000000..842348c1 --- /dev/null +++ b/geometric_kernels/utils/kernel_formulas/spd.py @@ -0,0 +1,127 @@ +""" +Implements an alternative formula for the heat kernel on the manifold of +symmetric positive definite matrices by :cite:t:`sawyer1992`. + +The implementation is adapted from https://github.com/imbirik/LieStationaryKernels. +Since the resulting approximation +* can fail to be positive semi-definite, +* is very slow, +* and is rather numerically unstable, +it is not recommended to use it in practice. The implementation is provided +mainly for testing purposes. +""" + +import lab as B +import numpy as np +import scipy +from beartype.typing import Optional + + +def _spd_heat_kernel_2x2_base( + t: float, + x: B.NPNumeric, + x2: Optional[B.NPNumeric] = None, +) -> float: + """ + The semi-analytic formula for the heat kernel on manifold of symmetric + positive definite matrices 2x2 from :cite:t:`sawyer1992`. The implementation + is adapted from https://github.com/imbirik/LieStationaryKernels. + + :param t: + The time parameter, a positive float. + :param x: + A single input, an array of shape [2, 2]. + :param x2: + A single input, an array of shape [2, 2]. If None, defaults to x. + + :return: + An approximation of the kernel value k(x, x2), a float. The kernel is not + normalized, i.e. k(x, x) may be an arbitrary (implementation-dependent) + positive number. For the normalized kernel which can also handle batch + inputs outputting covariance matrices, use :func:`sawyer_heat_kernel`. + """ + if x2 is None: + x2 = x + + assert x.shape == (2, 2) + assert x2.shape == (2, 2) + + cl_1 = np.linalg.cholesky(x) + cl_2 = np.linalg.cholesky(x2) + diff = np.linalg.inv(cl_2) @ cl_1 + _, singular_values, _ = np.linalg.svd(diff) + # Note: singular values that np.linalg.svd outputs are sorted, the following + # code relies on this fact. + H1, H2 = np.log(singular_values[0]), np.log(singular_values[1]) + assert H1 >= H2 + + r_H_sq = H1 * H1 + H2 * H2 + alpha = H1 - H2 + + # Non-integral part + result = 1.0 + result *= np.exp(-r_H_sq / (4 * t)) + + # Integrand + def link_function(x): + if x < 1e-5: + x = 1e-5 + res = 1.0 + res *= 2 * x + alpha + res *= np.exp(-x * (x + alpha) / (2 * t)) + res *= pow(np.sinh(x) * np.sinh(x + alpha), -1 / 2) + return res + + # Evaluating the integral + + # scipy.integrate.quad is much more accurate than np.trapz with + # b_vals = np.logspace(-3., 1, 1000), at least if we believe + # that Mathematica's NIntegrate is accurate. Also, you might think that + # scipy.integrate.quad_vec can be used to compute a whole covariance matrix + # at once. However, it seems to show terrible accuracy in this case. + + integral, error = scipy.integrate.quad(link_function, 0, np.inf) + + result *= integral + + return result + + +def spd_heat_kernel_2x2( + t: float, + X: B.NPNumeric, + X2: Optional[B.NPNumeric] = None, +) -> B.NPNumeric: + """ + The semi-analytic formula for the heat kernel on manifold of symmetric + positive definite matrices 2x2 from :cite:t:`sawyer1992`, normalized to + have k(x, x) = 1 for all x. The implementation is adapted from + https://github.com/imbirik/LieStationaryKernels. + + :param t: + The time parameter, a positive float. + :param X: + A batch of inputs, an array of shape [N, 2, 2]. + :param X2: + A batch of inputs, an array of shape [N2, 2, 2]. If None, defaults to X. + + :return: + The kernel matrix, an array of shape [N, N2]. The kernel is normalized, + i.e. k(x, x) = 1 for all x. + """ + + if X2 is None: + X2 = X + + normalization = _spd_heat_kernel_2x2_base(t, np.eye(2, 2)) + + result = np.zeros((X.shape[0], X2.shape[0])) + + # This is a very inefficient implementation, but it will do for tests. The + # straightforward vectorization of _sawyer_heat_kernel_base is not possible + # due to scipy.integrate.quad_vec giving very bad accuracy in this case. + for i, x in enumerate(X): + for j, x2 in enumerate(X2): + result[i, j] = _spd_heat_kernel_2x2_base(t, x, x2) / normalization + + return result diff --git a/geometric_kernels/utils/manifold_utils.py b/geometric_kernels/utils/manifold_utils.py index 8a0faeb4..42280cac 100644 --- a/geometric_kernels/utils/manifold_utils.py +++ b/geometric_kernels/utils/manifold_utils.py @@ -2,10 +2,89 @@ import lab as B import numpy as np +from beartype.typing import Optional from geometric_kernels.lab_extras import from_numpy +def minkowski_inner_product(vector_a: B.Numeric, vector_b: B.Numeric) -> B.Numeric: + r""" + Computes the Minkowski inner product of vectors. + + .. math:: \langle a, b \rangle = a_0 b_0 - a_1 b_1 - \ldots - a_n b_n. + + :param vector_a: + An [..., n+1]-shaped array of points in the hyperbolic space $\mathbb{H}_n$. + :param vector_b: + An [..., n+1]-shaped array of points in the hyperbolic space $\mathbb{H}_n$. + + :return: + An [...,]-shaped array of inner products. + """ + assert vector_a.shape == vector_b.shape + n = vector_a.shape[-1] - 1 + assert n > 0 + diagonal = from_numpy(vector_a, [-1.0] + [1.0] * n) # (n+1) + diagonal = B.cast(B.dtype(vector_a), diagonal) + return B.einsum("...i,...i->...", diagonal * vector_a, vector_b) + + +def hyperbolic_distance( + x1: B.Numeric, x2: B.Numeric, diag: Optional[bool] = False +) -> B.Numeric: + """ + Compute the hyperbolic distance between `x1` and `x2`. + + The code is a reimplementation of + `geomstats.geometry.hyperboloid.HyperbolicMetric` for `lab`. + + :param x1: + An [N, n+1]-shaped array of points in the hyperbolic space. + :param x2: + An [M, n+1]-shaped array of points in the hyperbolic space. + :param diag: + If True, compute elementwise distance. Requires N = M. + + Default False. + + :return: + An [N, M]-shaped array if diag=False or [N,]-shaped array + if diag=True. + """ + if diag: + # Compute a pointwise distance between `x1` and `x2` + x1_ = x1 + x2_ = x2 + else: + if B.rank(x1) == 1: + x1 = B.expand_dims(x1) + if B.rank(x2) == 1: + x2 = B.expand_dims(x2) + + # compute pairwise distance between arrays of points `x1` and `x2` + # `x1` (N, n+1) + # `x2` (M, n+1) + x1_ = B.tile(x1[..., None, :], 1, x2.shape[0], 1) # (N, M, n+1) + x2_ = B.tile(x2[None], x1.shape[0], 1, 1) # (N, M, n+1) + + sq_norm_1 = minkowski_inner_product(x1_, x1_) + sq_norm_2 = minkowski_inner_product(x2_, x2_) + inner_prod = minkowski_inner_product(x1_, x2_) + + cosh_angle = -inner_prod / B.sqrt(sq_norm_1 * sq_norm_2) + + one = B.cast(B.dtype(cosh_angle), from_numpy(cosh_angle, [1.0])) + large_constant = B.cast(B.dtype(cosh_angle), from_numpy(cosh_angle, [1e24])) + + # clip values into [1.0, 1e24] + cosh_angle = B.where(cosh_angle < one, one, cosh_angle) + cosh_angle = B.where(cosh_angle > large_constant, large_constant, cosh_angle) + + dist = B.log(cosh_angle + B.sqrt(cosh_angle**2 - 1)) # arccosh + dist = B.cast(B.dtype(x1_), dist) + return dist + + def manifold_laplacian(x: B.Numeric, manifold, egrad, ehess): r""" Computes the manifold Laplacian of a given function at a given point x. diff --git a/geometric_kernels/utils/product.py b/geometric_kernels/utils/product.py index 306d6a0b..58baded3 100644 --- a/geometric_kernels/utils/product.py +++ b/geometric_kernels/utils/product.py @@ -3,19 +3,30 @@ import lab as B from beartype.typing import Dict, List +from geometric_kernels.lab_extras import smart_cast -def params_to_params_list(params: Dict[str, B.Numeric]) -> List[Dict[str, B.Numeric]]: + +def params_to_params_list( + number_of_factors: int, params: Dict[str, B.Numeric] +) -> List[Dict[str, B.Numeric]]: """ Takes a dictionary of parameters of a product kernel and returns a list of - dictionaries of parameters for the factor kernels. + dictionaries of parameters for the factor kernels. The shape of "lengthscale" + should be the same as the shame of "nu", and the length of both should be + either 1 or equal to `number_of_factors`. + :param number_of_factors: + Number of factors in the product kernel. :param params: Parameters of the product kernel. """ assert params["lengthscale"].shape == params["nu"].shape assert len(params["nu"].shape) == 1 - number_of_factors = params["nu"].shape[0] + if params["nu"].shape[0] == 1: + return [params] * number_of_factors + + assert params["nu"].shape[0] == number_of_factors list_of_params: List[Dict[str, B.Numeric]] = [] for i in range(number_of_factors): @@ -43,12 +54,17 @@ def make_product(xs: List[B.Numeric]) -> B.Numeric: An [N, D]-shaped array, a batch of product space elements, where `D` is the sum, over all factor spaces, of `prod()`. """ - flat_xs = [B.reshape(x, B.shape(x)[0], -1) for x in xs] + common_dtype = B.promote_dtypes(*[B.dtype(x) for x in xs]) + + flat_xs = [B.cast(common_dtype, B.reshape(x, B.shape(x)[0], -1)) for x in xs] return B.concat(*flat_xs, axis=-1) def project_product( - x: B.Numeric, dimension_indices: List[List[int]], element_shapes: List[List[int]] + x: B.Numeric, + dimension_indices: List[List[int]], + element_shapes: List[List[int]], + element_dtypes: List[B.DType], ) -> List[B.Numeric]: """ Project an element of the product space onto each factor. @@ -63,7 +79,11 @@ def project_product( might be necessary to accommodate the spaces whose elements are matrices rather than vectors, as determined by `element_shapes`. :param element_shapes: - Shapes of the elements in each factor. + Shapes of the elements in each factor. Can be obtained as properties + `space.element_shape` of any given factor `space`. + :param element_dtypes: + Abstract lab data types of the elements in each factor. Can be obtained + as properties `space.element_dtype` of any given factor `space`. :return: A list of the batches of elements `xi` in factor spaces, each of the @@ -71,7 +91,7 @@ def project_product( """ N = x.shape[0] xs = [ - B.reshape(B.take(x, inds, axis=-1), N, *shape) - for inds, shape in zip(dimension_indices, element_shapes) + smart_cast(dtype, B.reshape(B.take(x, inds, axis=-1), N, *shape)) + for inds, shape, dtype in zip(dimension_indices, element_shapes, element_dtypes) ] return xs diff --git a/geometric_kernels/utils/special_functions.py b/geometric_kernels/utils/special_functions.py index 7e3b342f..b222191c 100644 --- a/geometric_kernels/utils/special_functions.py +++ b/geometric_kernels/utils/special_functions.py @@ -2,19 +2,15 @@ Special mathematical functions used in the library. """ -from math import sqrt - import lab as B from beartype.typing import List, Optional from geometric_kernels.lab_extras import ( count_nonzero, - float_like, from_numpy, int_like, take_along_axis, ) -from geometric_kernels.utils.utils import hamming_distance def walsh_function(d: int, combination: List[int], x: B.Bool) -> B.Float: @@ -113,42 +109,3 @@ def kravchuk_normalized( rhs_1 = (d - 2 * m) * kravchuk_normalized_j_minus_1 rhs_2 = -(j - 1) * kravchuk_normalized_j_minus_2 return (rhs_1 + rhs_2) / (d - j + 1) - - -def hypercube_graph_heat_kernel( - lengthscale: B.Numeric, - X: B.Numeric, - X2: Optional[B.Numeric] = None, - normalized_laplacian: bool = True, -): - """ - Analytic formula for the heat kernel on the hypercube graph, see - Equation (14) in :cite:t:`borovitskiy2023`. - - :param lengthscale: - The length scale of the kernel, an array of shape [1]. - :param X: - A batch of inputs, an array of shape [N, d]. - :param X2: - A batch of inputs, an array of shape [N2, d]. - - :return: - The kernel matrix, an array of shape [N, N2]. - """ - if X2 is None: - X2 = X - - assert lengthscale.shape == (1,) - assert X.ndim == 2 and X2.ndim == 2 - assert X.shape[-1] == X2.shape[-1] - - if normalized_laplacian: - d = X.shape[-1] - lengthscale = lengthscale / sqrt(d) - - # For TensorFlow, we need to explicitly cast the distances to double. - # Note: if we use B.dtype_float(X) instead of float_like(X), it gives - # float16 and TensorFlow is still complaining. - hamming_distances = B.cast(float_like(X), hamming_distance(X, X2)) - - return B.tanh(lengthscale**2 / 2) ** hamming_distances diff --git a/notebooks/Hyperbolic.ipynb b/notebooks/Hyperbolic.ipynb index f5757050..70e37da7 100644 --- a/notebooks/Hyperbolic.ipynb +++ b/notebooks/Hyperbolic.ipynb @@ -772,11 +772,14 @@ "```\n", "\n", "```\n", - "@article{azangulov2023,\n", - " title={Stationary Kernels and Gaussian Processes on Lie Groups and their Homogeneous Spaces II: non-compact symmetric spaces},\n", - " author={Azangulov, Iskander and Smolensky, Andrei and Terenin, Alexander and Borovitskiy, Viacheslav},\n", - " journal={arXiv preprint arXiv:2301.13088},\n", - " year={2023}\n", + "@article{azangulov2024b,\n", + " title = {Stationary Kernels and Gaussian Processes on Lie Groups and their Homogeneous Spaces II: non-compact symmetric spaces},\n", + " author = {Azangulov, Iskander and Smolensky, Andrei and Terenin, Alexander and Borovitskiy, Viacheslav},\n", + " journal = {Journal of Machine Learning Research},\n", + " year = {2024},\n", + " volume = {25},\n", + " number = {281},\n", + " pages = {1--51},\n", "}\n", "```" ] diff --git a/notebooks/SPD.ipynb b/notebooks/SPD.ipynb index 66e3911a..74a9255c 100644 --- a/notebooks/SPD.ipynb +++ b/notebooks/SPD.ipynb @@ -604,11 +604,14 @@ "```\n", "\n", "```\n", - "@article{azangulov2023,\n", - " title={Stationary Kernels and Gaussian Processes on Lie Groups and their Homogeneous Spaces II: non-compact symmetric spaces},\n", - " author={Azangulov, Iskander and Smolensky, Andrei and Terenin, Alexander and Borovitskiy, Viacheslav},\n", - " journal={arXiv preprint arXiv:2301.13088},\n", - " year={2023}\n", + "@article{azangulov2024b,\n", + " title = {Stationary Kernels and Gaussian Processes on Lie Groups and their Homogeneous Spaces II: non-compact symmetric spaces},\n", + " author = {Azangulov, Iskander and Smolensky, Andrei and Terenin, Alexander and Borovitskiy, Viacheslav},\n", + " journal = {Journal of Machine Learning Research},\n", + " year = {2024},\n", + " volume = {25},\n", + " number = {281},\n", + " pages = {1--51},\n", "}\n", "```" ] diff --git a/notebooks/SpecialOrthogonal.ipynb b/notebooks/SpecialOrthogonal.ipynb index a8f66ac3..c2a47e53 100644 --- a/notebooks/SpecialOrthogonal.ipynb +++ b/notebooks/SpecialOrthogonal.ipynb @@ -824,11 +824,14 @@ "```\n", "\n", "```\n", - "@article{azangulov2022,\n", - " title={Stationary kernels and Gaussian processes on Lie groups and their homogeneous spaces I: the compact case},\n", - " author={Azangulov, Iskander and Smolensky, Andrei and Terenin, Alexander and Borovitskiy, Viacheslav},\n", - " journal={arXiv preprint arXiv:2208.14960},\n", - " year={2022}\n", + "@article{azangulov2024a,\n", + " title = {Stationary Kernels and Gaussian Processes on Lie Groups and their Homogeneous Spaces I: the compact case},\n", + " author = {Azangulov, Iskander and Smolensky, Andrei and Terenin, Alexander and Borovitskiy, Viacheslav},\n", + " journal = {Journal of Machine Learning Research},\n", + " year = {2024},\n", + " volume = {25},\n", + " number = {280},\n", + " pages = {1--52},\n", "}\n", "```" ] diff --git a/notebooks/SpecialUnitary.ipynb b/notebooks/SpecialUnitary.ipynb index 45ed5320..0c0f2165 100644 --- a/notebooks/SpecialUnitary.ipynb +++ b/notebooks/SpecialUnitary.ipynb @@ -882,11 +882,14 @@ "```\n", "\n", "```\n", - "@article{azangulov2022,\n", - " title={Stationary kernels and Gaussian processes on Lie groups and their homogeneous spaces I: the compact case},\n", - " author={Azangulov, Iskander and Smolensky, Andrei and Terenin, Alexander and Borovitskiy, Viacheslav},\n", - " journal={arXiv preprint arXiv:2208.14960},\n", - " year={2022}\n", + "@article{azangulov2024a,\n", + " title = {Stationary Kernels and Gaussian Processes on Lie Groups and their Homogeneous Spaces I: the compact case},\n", + " author = {Azangulov, Iskander and Smolensky, Andrei and Terenin, Alexander and Borovitskiy, Viacheslav},\n", + " journal = {Journal of Machine Learning Research},\n", + " year = {2024},\n", + " volume = {25},\n", + " number = {280},\n", + " pages = {1--52},\n", "}\n", "```" ] diff --git a/notebooks/other/Hyperbolic Approximations.ipynb b/notebooks/other/Hyperbolic Approximations.ipynb index c8f128eb..669d1ea7 100644 --- a/notebooks/other/Hyperbolic Approximations.ipynb +++ b/notebooks/other/Hyperbolic Approximations.ipynb @@ -659,11 +659,14 @@ "```\n", "\n", "```\n", - "@article{azangulov2023,\n", - " title={Stationary Kernels and Gaussian Processes on Lie Groups and their Homogeneous Spaces II: non-compact symmetric spaces},\n", - " author={Azangulov, Iskander and Smolensky, Andrei and Terenin, Alexander and Borovitskiy, Viacheslav},\n", - " journal={arXiv preprint arXiv:2301.13088},\n", - " year={2023}\n", + "@article{azangulov2024b,\n", + " title = {Stationary Kernels and Gaussian Processes on Lie Groups and their Homogeneous Spaces II: non-compact symmetric spaces},\n", + " author = {Azangulov, Iskander and Smolensky, Andrei and Terenin, Alexander and Borovitskiy, Viacheslav},\n", + " journal = {Journal of Machine Learning Research},\n", + " year = {2024},\n", + " volume = {25},\n", + " number = {281},\n", + " pages = {1--51},\n", "}\n", "```" ] diff --git a/notebooks/other/SPD Approximations.ipynb b/notebooks/other/SPD Approximations.ipynb index ae1e5b14..7736a729 100644 --- a/notebooks/other/SPD Approximations.ipynb +++ b/notebooks/other/SPD Approximations.ipynb @@ -530,11 +530,14 @@ "```\n", "\n", "```\n", - "@article{azangulov2023,\n", - " title={Stationary Kernels and Gaussian Processes on Lie Groups and their Homogeneous Spaces II: non-compact symmetric spaces},\n", - " author={Azangulov, Iskander and Smolensky, Andrei and Terenin, Alexander and Borovitskiy, Viacheslav},\n", - " journal={arXiv preprint arXiv:2301.13088},\n", - " year={2023}\n", + "@article{azangulov2024b,\n", + " title = {Stationary Kernels and Gaussian Processes on Lie Groups and their Homogeneous Spaces II: non-compact symmetric spaces},\n", + " author = {Azangulov, Iskander and Smolensky, Andrei and Terenin, Alexander and Borovitskiy, Viacheslav},\n", + " journal = {Journal of Machine Learning Research},\n", + " year = {2024},\n", + " volume = {25},\n", + " number = {281},\n", + " pages = {1--51},\n", "}\n", "```" ] diff --git a/geometric_kernels/utils/compute_characters.py b/scripts/compute_characters.py similarity index 99% rename from geometric_kernels/utils/compute_characters.py rename to scripts/compute_characters.py index 9b4d4aa3..412b7405 100644 --- a/geometric_kernels/utils/compute_characters.py +++ b/scripts/compute_characters.py @@ -141,7 +141,7 @@ def indent_str(self) -> str: def compute_character_formula_so(self, signature): """ - Refer to the appendix of cite:t:`azangulov2022`, + Refer to the appendix of cite:t:`azangulov2024a`, https://arxiv.org/pdf/2208.14960.pdf. """ n = self.n @@ -259,7 +259,7 @@ def xi1(qs): def compute_character_formula_su(self, signature): """ - Refer to the appendix of cite:t:`azangulov2022`, + Refer to the appendix of cite:t:`azangulov2024a`, https://arxiv.org/pdf/2208.14960.pdf. """ n = self.n diff --git a/test_requirements.txt b/test_requirements.txt index 56873988..e81ad77a 100644 --- a/test_requirements.txt +++ b/test_requirements.txt @@ -12,7 +12,7 @@ flake8==7.0.0 isort==5.13.2 autoflake pytest -# pytest-cov +pytest-cov nbmake mypy diff --git a/tests/data.py b/tests/data.py new file mode 100644 index 00000000..b267e894 --- /dev/null +++ b/tests/data.py @@ -0,0 +1,47 @@ +from math import sqrt as sr +from pathlib import Path + +import numpy as np + +TEST_MESH_PATH = str(Path(__file__).parent.resolve() / "teddy.obj") + +TEST_GRAPH_ADJACENCY = np.array( + [ + [0, 1, 0, 0, 0, 0, 0], + [1, 0, 1, 1, 1, 0, 0], + [0, 1, 0, 0, 0, 1, 0], + [0, 1, 0, 0, 1, 0, 0], + [0, 1, 0, 1, 0, 0, 0], + [0, 0, 1, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0, 0], + ] +).astype(np.float64) + + +TEST_GRAPH_LAPLACIAN = np.array( + [ + [1, -1, 0, 0, 0, 0, 0], + [-1, 4, -1, -1, -1, 0, 0], + [0, -1, 2, 0, 0, -1, 0], + [0, -1, 0, 2, -1, 0, 0], + [0, -1, 0, -1, 2, 0, 0], + [0, 0, -1, 0, 0, 1, 0], + [0, 0, 0, 0, 0, 0, 0], + ] +).astype( + np.float64 +) # corresponds to TEST_GRAPH_ADJACENCY, unnormalized Laplacian + +TEST_GRAPH_LAPLACIAN_NORMALIZED = np.array( + [ + [1, -0.5, 0, 0, 0, 0, 0], # noqa: E241 + [-0.5, 1, -1 / sr(2) / 2, -1 / sr(2) / 2, -1 / sr(2) / 2, 0, 0], # noqa: E241 + [0, -1 / sr(2) / 2, 1, 0, 0, -1 / sr(2), 0], # noqa: E241 + [0, -1 / sr(2) / 2, 0, 1, -0.5, 0, 0], # noqa: E241 + [0, -1 / sr(2) / 2, 0, -0.5, 1, 0, 0], # noqa: E241 + [0, 0, -1 / sr(2), 0, 0, 1, 0], # noqa: E241 + [0, 0, 0, 0, 0, 0, 0], # noqa: E241 + ] +).astype( + np.float64 +) # corresponds to TEST_GRAPH_ADJACENCY, normalized Laplacian diff --git a/tests/feature_maps/__init__.py b/tests/feature_maps/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/feature_maps/test_feature_maps.py b/tests/feature_maps/test_feature_maps.py new file mode 100644 index 00000000..db7f6c90 --- /dev/null +++ b/tests/feature_maps/test_feature_maps.py @@ -0,0 +1,76 @@ +import lab as B +import numpy as np +import pytest + +from geometric_kernels.feature_maps import RandomPhaseFeatureMapCompact +from geometric_kernels.kernels import MaternGeometricKernel, default_feature_map +from geometric_kernels.kernels.matern_kernel import default_num +from geometric_kernels.spaces import NoncompactSymmetricSpace +from geometric_kernels.utils.utils import make_deterministic + +from ..helper import check_function_with_backend, create_random_state, spaces + + +@pytest.fixture( + params=spaces(), + ids=str, +) +def feature_map_and_friends(request, backend): + """ + Returns a tuple (feature_map, kernel, space) where: + - feature_map is the `default_feature_map` of the `kernel`, + - kernel is the `MaternGeometricKernel` on the `space`, with a reasonably + small value of `num`, + - space = request.param, + + `backend` parameter is required to create a random state for the feature + map, if it requires one. + """ + space = request.param + + if isinstance(space, NoncompactSymmetricSpace): + kernel = MaternGeometricKernel( + space, key=create_random_state(backend), num=min(default_num(space), 100) + ) + else: + kernel = MaternGeometricKernel(space, num=min(default_num(space), 3)) + + feature_map = default_feature_map(kernel=kernel) + if isinstance(feature_map, RandomPhaseFeatureMapCompact): + # RandomPhaseFeatureMapCompact requires a key. Note: normally, + # RandomPhaseFeatureMapNoncompact, RejectionSamplingFeatureMapHyperbolic, + # and RejectionSamplingFeatureMapSPD also require a key, but when they + # are obtained from an already constructed kernel's feature map, the key + # is already provided and fixed in the similar way as we do just below. + feature_map = make_deterministic(feature_map, key=create_random_state(backend)) + + return feature_map, kernel, space + + +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_feature_map_approximates_kernel(backend, feature_map_and_friends): + feature_map, kernel, space = feature_map_and_friends + + params = kernel.init_params() + + key = np.random.RandomState(0) + key, X = space.random(key, 50) + + def diff_kern_mats(params, X): + _, embedding = feature_map(X, params) + + kernel_mat = kernel.K(params, X, X) + kernel_mat_alt = B.matmul(embedding, B.T(embedding)) + + return kernel_mat - kernel_mat_alt + + # Check that, approximately, k(X, X) = , where k is the + # kernel and phi is the feature map. + check_function_with_backend( + backend, + np.zeros((X.shape[0], X.shape[0])), + diff_kern_mats, + params, + X, + atol=0.1, + ) diff --git a/tests/sampling/test_student_t_sample.py b/tests/feature_maps/test_student_t_sample.py similarity index 95% rename from tests/sampling/test_student_t_sample.py rename to tests/feature_maps/test_student_t_sample.py index 5b262f15..9a1051d7 100644 --- a/tests/sampling/test_student_t_sample.py +++ b/tests/feature_maps/test_student_t_sample.py @@ -9,7 +9,7 @@ def test_student_t_sample(deg_freedom, n): size = (2048,) - key = np.random.RandomState(seed=1234) + key = np.random.RandomState(0) shape = 1.0 * np.eye(n) loc = 1.0 * np.zeros((n,)) diff --git a/tests/helper.py b/tests/helper.py index 3dad2210..9c3613ad 100644 --- a/tests/helper.py +++ b/tests/helper.py @@ -1,12 +1,87 @@ import lab as B import numpy as np from beartype.door import die_if_unbearable, is_bearable -from beartype.typing import Any, Callable, Optional, Union +from beartype.typing import Any, Callable, List, Optional, Union from plum import ModuleType, resolve_type_hint +from geometric_kernels.lab_extras import SparseArray +from geometric_kernels.spaces import ( + Circle, + CompactMatrixLieGroup, + DiscreteSpectrumSpace, + Graph, + Hyperbolic, + HypercubeGraph, + Hypersphere, + Mesh, + NoncompactSymmetricSpace, + ProductDiscreteSpectrumSpace, + Space, + SpecialOrthogonal, + SpecialUnitary, + SymmetricPositiveDefiniteMatrices, +) + +from .data import TEST_GRAPH_ADJACENCY, TEST_MESH_PATH + EagerTensor = ModuleType("tensorflow.python.framework.ops", "EagerTensor") +def compact_matrix_lie_groups() -> List[CompactMatrixLieGroup]: + return [ + SpecialOrthogonal(3), + SpecialOrthogonal(8), + SpecialUnitary(2), + SpecialUnitary(5), + ] + + +def product_discrete_spectrum_spaces() -> List[ProductDiscreteSpectrumSpace]: + return [ + ProductDiscreteSpectrumSpace(Circle(), Hypersphere(3), Circle()), + ProductDiscreteSpectrumSpace( + Circle(), Graph(np.kron(TEST_GRAPH_ADJACENCY, TEST_GRAPH_ADJACENCY)) + ), # TEST_GRAPH_ADJACENCY is too small for default parameters of the ProductDiscreteSpectrumSpace + ProductDiscreteSpectrumSpace(Mesh.load_mesh(TEST_MESH_PATH), Hypersphere(2)), + ] + + +def discrete_spectrum_spaces() -> List[DiscreteSpectrumSpace]: + return ( + [ + Circle(), + HypercubeGraph(1), + HypercubeGraph(3), + HypercubeGraph(6), + Hypersphere(2), + Hypersphere(3), + Hypersphere(10), + Mesh.load_mesh(TEST_MESH_PATH), + Graph(TEST_GRAPH_ADJACENCY, normalize_laplacian=False), + Graph(TEST_GRAPH_ADJACENCY, normalize_laplacian=True), + ] + + compact_matrix_lie_groups() + + product_discrete_spectrum_spaces() + ) + + +def noncompact_symmetric_spaces() -> List[NoncompactSymmetricSpace]: + return [ + Hyperbolic(2), + Hyperbolic(3), + Hyperbolic(8), + Hyperbolic(9), + SymmetricPositiveDefiniteMatrices(2), + SymmetricPositiveDefiniteMatrices(3), + SymmetricPositiveDefiniteMatrices(6), + SymmetricPositiveDefiniteMatrices(7), + ] + + +def spaces() -> List[Space]: + return discrete_spectrum_spaces() + noncompact_symmetric_spaces() + + def np_to_backend(value: B.NPNumeric, backend: str): """ Converts a numpy array to the desired backend. @@ -36,16 +111,26 @@ def np_to_backend(value: B.NPNumeric, backend: str): import jax.numpy as jnp return jnp.array(value) + elif backend == "scipy_sparse": + import scipy.sparse as sp + + return sp.csr_array(value) else: raise ValueError("Unknown backend: {}".format(backend)) +def create_random_state(backend: str, seed: int = 0): + dtype = B.dtype(np_to_backend(np.array([1.0]), backend)) + return B.create_random_state(dtype, seed=seed) + + def array_type(backend: str): """ Returns the array type corresponding to the given backend. :param backend: - The backend to use, one of the strings "tensorflow", "torch", "numpy", "jax". + The backend to use, one of the strings "tensorflow", "torch", "numpy", + "jax", "scipy_sparse". :return: The array type corresponding to the given backend. @@ -58,10 +143,33 @@ def array_type(backend: str): return resolve_type_hint(B.NPNumeric) elif backend == "jax": return resolve_type_hint(B.JAXNumeric) + elif backend == "scipy_sparse": + return resolve_type_hint(SparseArray) else: raise ValueError(f"Unknown backend: {backend}") +def apply_recursive(data: Any, func: Callable[[Any], Any]) -> Any: + """ + Apply a function recursively to a nested data structure. Supports lists and + dictionaries. + + :param data: + The data structure to apply the function to. + :param func: + The function to apply. + + :return: + The data structure with the function applied to each element. + """ + if isinstance(data, dict): + return {key: apply_recursive(value, func) for key, value in data.items()} + elif isinstance(data, list): + return [apply_recursive(element, func) for element in data] + else: + return func(data) + + def check_function_with_backend( backend: str, result: Any, @@ -97,17 +205,31 @@ def check_function_with_backend( the expected result. """ - args_casted = [] - for arg in args: + def cast(arg): if is_bearable(arg, B.Numeric): # We only expect numpy arrays here die_if_unbearable(arg, B.NPNumeric) - args_casted.append(np_to_backend(arg, backend)) + return np_to_backend(arg, backend) else: - args_casted.append(arg) + return arg + + args_casted = (apply_recursive(arg, cast) for arg in args) + f_output = f(*args_casted) - assert is_bearable(f_output, array_type(backend)) + assert is_bearable( + f_output, array_type(backend) + ), f"The output is not of the expected type. Expected: {array_type(backend)}, got: {type(f_output)}" if compare_to_result is None: - np.testing.assert_allclose(B.to_numpy(f_output), result, atol=atol) + # we convert `f_output` to numpy array to compare with `result`` + if is_bearable(f_output, SparseArray): + f_output = f_output.toarray() + else: + f_output = B.to_numpy(f_output) + assert ( + result.shape == f_output.shape + ), f"Shapes do not match: {result.shape} (for result) vs {f_output.shape} (for f_output)" + np.testing.assert_allclose(f_output, result, atol=atol) else: - assert compare_to_result(result, f_output) + assert compare_to_result( + result, f_output + ), f"compare_to_result(result, f_output) failed with\n result:\n{result}\n\nf_output:\n{f_output}" diff --git a/tests/kernels/test_feature_map_kernel.py b/tests/kernels/test_feature_map_kernel.py new file mode 100644 index 00000000..d87a289a --- /dev/null +++ b/tests/kernels/test_feature_map_kernel.py @@ -0,0 +1,123 @@ +import lab as B +import numpy as np +import pytest + +from geometric_kernels.kernels import MaternFeatureMapKernel +from geometric_kernels.kernels.matern_kernel import default_feature_map, default_num + +from ..helper import ( + check_function_with_backend, + create_random_state, + noncompact_symmetric_spaces, +) + + +@pytest.fixture( + params=noncompact_symmetric_spaces(), + ids=str, + scope="module", +) +def inputs(request): + """ + Returns a tuple (space, num_features, feature_map, X, X2) where: + - space = request.param, + - num_features = default_num(space) or 15, whichever is smaller, + - feature_map = default_feature_map(space=space, num=num_features), + - X is a random sample of random size from the space, + - X2 is another random sample of random size from the space, + """ + space = request.param + num_features = min(default_num(space), 15) + feature_map = default_feature_map(space=space, num=num_features) + + key = np.random.RandomState(0) + N, N2 = key.randint(low=1, high=100 + 1, size=2) + key, X = space.random(key, N) + key, X2 = space.random(key, N2) + + return space, num_features, feature_map, X, X2 + + +@pytest.fixture +def kernel(inputs, backend, normalize=True): + space, _, feature_map, _, _ = inputs + + key = create_random_state(backend) + + return MaternFeatureMapKernel(space, feature_map, key, normalize=normalize) + + +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_params(inputs, backend, kernel): + params = kernel.init_params() + + assert "lengthscale" in params + assert params["lengthscale"].shape == (1,) + assert "nu" in params + assert params["nu"].shape == (1,) + + +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +@pytest.mark.parametrize("normalize", [True, False], ids=["normalize", "no_normalize"]) +def test_K(inputs, backend, normalize, kernel): + _, _, _, X, X2 = inputs + params = kernel.init_params() + + # Check that kernel.K runs and the output is a tensor of the right backend and shape. + check_function_with_backend( + backend, + (X.shape[0], X2.shape[0]), + kernel.K, + params, + X, + X2, + compare_to_result=lambda res, f_out: B.shape(f_out) == res, + ) + + +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +@pytest.mark.parametrize("normalize", [True, False], ids=["normalize", "no_normalize"]) +def test_K_one_param(inputs, backend, normalize, kernel): + _, _, _, X, _ = inputs + params = kernel.init_params() + + # Check that kernel.K(X) coincides with kernel.K(X, X). + check_function_with_backend( + backend, + np.zeros((X.shape[0], X.shape[0])), + lambda params, X: kernel.K(params, X) - kernel.K(params, X, X), + params, + X, + ) + + +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +@pytest.mark.parametrize("normalize", [True, False], ids=["normalize", "no_normalize"]) +def test_K_diag(inputs, backend, normalize, kernel): + _, _, _, X, _ = inputs + params = kernel.init_params() + + # Check that kernel.K_diag coincides with the diagonal of kernel.K. + check_function_with_backend( + backend, + np.zeros((X.shape[0],)), + lambda params, X: kernel.K_diag(params, X) - B.diag(kernel.K(params, X)), + params, + X, + ) + + +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_normalize(inputs, backend, kernel): + _, _, _, X, _ = inputs + + params = kernel.init_params() + + # Check that the variance of the kernel is constant 1. + check_function_with_backend( + backend, + np.ones((X.shape[0],)), + kernel.K_diag, + params, + X, + ) diff --git a/tests/kernels/test_feature_maps.py b/tests/kernels/test_feature_maps.py deleted file mode 100644 index 926e751a..00000000 --- a/tests/kernels/test_feature_maps.py +++ /dev/null @@ -1,69 +0,0 @@ -from pathlib import Path - -import numpy as np -import pytest - -from geometric_kernels.kernels import MaternGeometricKernel, default_feature_map -from geometric_kernels.spaces.circle import Circle -from geometric_kernels.spaces.graph import Graph -from geometric_kernels.spaces.hyperbolic import Hyperbolic -from geometric_kernels.spaces.hypersphere import Hypersphere -from geometric_kernels.spaces.mesh import Mesh -from geometric_kernels.spaces.spd import SymmetricPositiveDefiniteMatrices - - -@pytest.mark.parametrize( - "space_name", ["circle", "hypersphere", "mesh", "graph", "hyperbolic", "spd"] -) -def test_feature_maps(space_name): - key = np.random.RandomState(1234) - num_points = 5 - - if space_name == "circle": - space = Circle() - kwargs = {} - elif space_name == "hypersphere": - space = Hypersphere(2) - kwargs = {} - elif space_name == "mesh": - filename = Path(__file__).parent / "../teddy.obj" - space = Mesh.load_mesh(str(filename)) - kwargs = {} - elif space_name == "graph": - A = np.array( - [ - [0, 1, 0, 0, 0, 0, 0], - [1, 0, 1, 1, 1, 0, 0], - [0, 1, 0, 0, 0, 1, 0], - [0, 1, 0, 0, 1, 0, 0], - [0, 1, 0, 1, 0, 0, 0], - [0, 0, 1, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0], - ] - ).astype("float") - - space = Graph(A, normalize_laplacian=True) - kwargs = {} - elif space_name == "hyperbolic": - space = Hyperbolic(dim=2) - kwargs = {"key": key} - elif space_name == "spd": - space = SymmetricPositiveDefiniteMatrices(n=2) - kwargs = {"key": key} - else: - raise ValueError(f"Unknown space {space}") - - kernel = MaternGeometricKernel(space, **kwargs) - params = kernel.init_params() - params = {"nu": np.r_[2.5], "lengthscale": np.r_[1.0]} - - feature_map = default_feature_map(kernel=kernel) - - key, points = space.random(key, num_points) - - _, embedding = feature_map(points, params, **kwargs) - - kernel_mat = kernel.K(params, points, points) - kernel_mat_alt = np.matmul(embedding, embedding.T) - - np.testing.assert_allclose(kernel_mat, kernel_mat_alt) diff --git a/tests/kernels/test_matern_karhunenloeve_kernel.py b/tests/kernels/test_matern_karhunenloeve_kernel.py index 1d87ecb2..802d6e61 100644 --- a/tests/kernels/test_matern_karhunenloeve_kernel.py +++ b/tests/kernels/test_matern_karhunenloeve_kernel.py @@ -1,62 +1,192 @@ -from pathlib import Path +from itertools import product +import lab as B import numpy as np -from pytest import fixture +import pytest from geometric_kernels.kernels import MaternKarhunenLoeveKernel +from geometric_kernels.kernels.matern_kernel import default_num from geometric_kernels.spaces import Mesh -_TRUNCATION_LEVEL = 10 -_NU = np.r_[1 / 2.0] +from ..helper import check_function_with_backend, discrete_spectrum_spaces +_EPS = 1e-5 -@fixture(name="kernel") -def fixture_mesh_kernel() -> MaternKarhunenLoeveKernel: - filename = Path(__file__).parent / "../teddy.obj" - mesh = Mesh.load_mesh(str(filename)) - return MaternKarhunenLoeveKernel(mesh, _TRUNCATION_LEVEL) +@pytest.fixture( + params=product(discrete_spectrum_spaces(), [True, False]), + ids=lambda tpl: f"{tpl[0]}{'-normalized' if tpl[1] else ''}", + scope="module", +) +def inputs(request): + """ + Returns a tuple (space, num_levels, kernel, X, X2) where: + - space = request.param[0], + - num_levels = default_num(space), + - kernel = MaternKarhunenLoeveKernel(space, num_levels, normalize=request.param[1]), + - X is a random sample of random size from the space, + - X2 is another random sample of random size from the space, + """ + space, normalize = request.param + num_levels = default_num(space) + kernel = MaternKarhunenLoeveKernel(space, num_levels, normalize=normalize) -def test_eigenvalues(kernel: MaternKarhunenLoeveKernel): - params = dict(lengthscale=np.r_[0.81], nu=_NU) + key = np.random.RandomState(0) + N, N2 = key.randint(low=1, high=100 + 1, size=2) + key, X = space.random(key, N) + key, X2 = space.random(key, N2) - assert kernel.eigenvalues(params).shape == (_TRUNCATION_LEVEL, 1) + return space, num_levels, kernel, X, X2 -def test_eigenfunctions(kernel: MaternKarhunenLoeveKernel): - num_data = 11 - Phi = kernel.eigenfunctions - X = np.random.randint(low=0, high=kernel.space.num_vertices, size=(num_data, 1)) +def test_params(inputs): + _, _, kernel, _, _ = inputs - assert Phi(X, lengthscale=np.r_[0.93]).shape == (num_data, _TRUNCATION_LEVEL) + params = kernel.init_params() + + assert "lengthscale" in params + assert params["lengthscale"].shape == (1,) + assert "nu" in params + assert params["nu"].shape == (1,) + + +def test_num_levels(inputs): + _, num_levels, kernel, _, _ = inputs + + assert kernel.eigenfunctions.num_levels == num_levels + + +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_eigenvalues_shape(inputs, backend): + _, num_levels, kernel, _, _ = inputs + params = kernel.init_params() + + # Check that the eigenvalues have appropriate shape. + check_function_with_backend( + backend, + (num_levels, 1), + kernel.eigenvalues, + params, + compare_to_result=lambda res, f_out: B.shape(f_out) == res, + ) + + +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_eigenvalues_positive(inputs, backend): + _, _, kernel, _, _ = inputs + params = kernel.init_params() + # Check that the eigenvalues are nonnegative. + check_function_with_backend( + backend, + None, + kernel.eigenvalues, + params, + compare_to_result=lambda _, f_out: np.all(B.to_numpy(f_out) >= 0), + ) -def test_K_shapes(kernel: MaternKarhunenLoeveKernel): + +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_eigenvalues_ordered(inputs, backend): + _, _, kernel, _, _ = inputs + params = kernel.init_params() + + # Check that the eigenvalues are sorted in descending order. + check_function_with_backend( + backend, + None, + kernel.eigenvalues, + params, + compare_to_result=lambda _, f_out: np.all( + B.to_numpy(f_out)[:-1] >= B.to_numpy(f_out)[1:] - _EPS + ), + ) + + +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_K(inputs, backend): + _, _, kernel, X, X2 = inputs params = kernel.init_params() - params["nu"] = _NU - params["lengthscale"] = np.r_[0.99] - N1, N2 = 11, 13 - X = np.random.randint(low=0, high=kernel.space.num_vertices, size=(N1, 1)) - X2 = np.random.randint(low=0, high=kernel.space.num_vertices, size=(N2, 1)) + result = kernel.K(params, X, X2) - K = kernel.K(params, X, None) - assert K.shape == (N1, N1) + assert result.shape == (X.shape[0], X2.shape[0]), "K has incorrect shape" - K = kernel.K(params, X, X2) - assert K.shape == (N1, N2) + if backend != "numpy": + # Check that kernel.K computed using `backend` coincides with the numpy result. + check_function_with_backend( + backend, + result, + kernel.K, + params, + X, + X2, + ) + + +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_K_one_param(inputs, backend): + space, _, kernel, X, _ = inputs + params = kernel.init_params() - K = kernel.K_diag(params, X) - assert K.shape == (N1,) + result = kernel.K(params, X, X) + # Check that kernel.K(X) coincides with kernel.K(X, X). + check_function_with_backend( + backend, + result, + kernel.K, + params, + X, + atol=1e-2 if isinstance(space, Mesh) else _EPS, + ) -def test_normalize(kernel: MaternKarhunenLoeveKernel): - random_points = np.arange(kernel.space.num_vertices).reshape(-1, 1) +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_K_diag(inputs, backend): + space, _, kernel, X, _ = inputs params = kernel.init_params() - params["nu"] = _NU - params["lengthscale"] = np.r_[0.99] - K = kernel.K_diag(params, random_points) # (N, ) + result = kernel.K(params, X).diagonal() - np.testing.assert_allclose(np.mean(K), 1.0) + assert result.shape == (X.shape[0],), "The diagonal has incorrect shape" + + # Check that kernel.K_diag coincides with the diagonal of kernel.K. + check_function_with_backend( + backend, + result, + kernel.K_diag, + params, + X, + atol=1e-2 if isinstance(space, Mesh) else _EPS, + ) + + +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_normalize(inputs, backend): + space, _, kernel, _, _ = inputs + + if not kernel.normalize: + pytest.skip("No need to check normalization for an unnormalized kernel") + + params = kernel.init_params() + key = np.random.RandomState(0) + key, X = space.random( + key, 1000 + ) # we need a large sample to get a good estimate of the mean variance + + def mean_variance(params, X): + return B.reshape( + B.mean(kernel.K_diag(params, X), squeeze=False), + 1, + ) # the reshape shields from a bug in lab present at least up to version 1.6.6 + + # Check that the average variance of the kernel is 1. + check_function_with_backend( + backend, + np.array([1.0]), + mean_variance, + params, + X, + atol=0.2, # very loose, but helps make sure the result is close to 1 + ) diff --git a/tests/kernels/test_normalization.py b/tests/kernels/test_normalization.py deleted file mode 100644 index c95d02a3..00000000 --- a/tests/kernels/test_normalization.py +++ /dev/null @@ -1,87 +0,0 @@ -from pathlib import Path - -import numpy as np -import pytest - -from geometric_kernels.feature_maps import RandomPhaseFeatureMapNoncompact -from geometric_kernels.kernels import MaternFeatureMapKernel, MaternKarhunenLoeveKernel -from geometric_kernels.spaces.circle import Circle -from geometric_kernels.spaces.graph import Graph -from geometric_kernels.spaces.hyperbolic import Hyperbolic -from geometric_kernels.spaces.hypersphere import Hypersphere -from geometric_kernels.spaces.mesh import Mesh -from geometric_kernels.spaces.spd import SymmetricPositiveDefiniteMatrices - - -@pytest.mark.parametrize("space_name", ["circle", "hypersphere", "mesh", "graph"]) -def test_normalization_matern_kl_kernel(space_name): - key = np.random.RandomState(1234) - num_points = 300 - num_eigenfns = 10 - - if space_name == "circle": - space = Circle() - key, points = space.random(key, num_points) - elif space_name == "hypersphere": - space = Hypersphere(2) - key, points = space.random(key, num_points) - elif space_name == "mesh": - filename = Path(__file__).parent / "../teddy.obj" - space = Mesh.load_mesh(str(filename)) - points = np.arange(space.num_vertices).reshape(-1, 1) - elif space_name == "graph": - A = np.array( - [ - [0, 1, 0, 0, 0, 0, 0], - [1, 0, 1, 1, 1, 0, 0], - [0, 1, 0, 0, 0, 1, 0], - [0, 1, 0, 0, 1, 0, 0], - [0, 1, 0, 1, 0, 0, 0], - [0, 0, 1, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0], - ] - ).astype("float") - - space = Graph(A, normalize_laplacian=True) - points = np.arange(space.num_vertices).reshape(-1, 1) - - else: - raise ValueError(f"Unknown space {space}") - - kernel = MaternKarhunenLoeveKernel(space, num_eigenfns, normalize=True) - params = kernel.init_params() - params = {"nu": np.r_[2.5], "lengthscale": np.r_[1.0]} - - kxx = kernel.K_diag(params, points) - np.testing.assert_allclose(np.mean(kxx), 1.0) - - kxx = np.diag(kernel.K(params, points)) - np.testing.assert_allclose(np.mean(kxx), 1.0) - - -@pytest.mark.parametrize("space_name", ["hyperbolic", "spd"]) -def test_normalization_feature_map_kernel(space_name): - key = np.random.RandomState(1234) - num_points = 300 - num_features = 10 - - if space_name == "hyperbolic": - space = Hyperbolic(dim=2) - points = space.random_point(num_points) - elif space_name == "spd": - space = SymmetricPositiveDefiniteMatrices(n=2) - points = space.random_point(num_points) - else: - raise ValueError(f"Unknown space {space}") - - params = dict(nu=np.r_[2.5], lengthscale=np.r_[1.0]) - - feature_map = RandomPhaseFeatureMapNoncompact(space, num_features) - - kernel = MaternFeatureMapKernel(space, feature_map, key) - - kxx = kernel.K_diag(params, points) - np.testing.assert_allclose(kxx, 1.0) - - kxx = np.diag(kernel.K(params, points)) - np.testing.assert_allclose(kxx, 1.0) diff --git a/tests/kernels/test_product.py b/tests/kernels/test_product.py deleted file mode 100644 index bb345c37..00000000 --- a/tests/kernels/test_product.py +++ /dev/null @@ -1,189 +0,0 @@ -import lab as B -import numpy as np - -from geometric_kernels.kernels import MaternKarhunenLoeveKernel, ProductGeometricKernel -from geometric_kernels.lab_extras.extras import from_numpy -from geometric_kernels.spaces import ( - Circle, - Hypersphere, - ProductDiscreteSpectrumSpace, - SpecialUnitary, -) -from geometric_kernels.utils.product import make_product -from geometric_kernels.utils.utils import chain - -_TRUNC_LEVEL = 128 -_GRID_SIZE = 3 - - -def test_circle_product_eigenfunctions(): - # assert that the naive method of phi-product calculation - # gives the same result as the addition theorem based calculation - product = ProductDiscreteSpectrumSpace( - Circle(), Circle(), num_levels=_TRUNC_LEVEL**2 - ) - - grid = B.linspace(0, 2 * B.pi, _GRID_SIZE) - ones = B.ones(_GRID_SIZE) - grid = B.stack( - grid[:, None] * ones[None, :], grid[None, :] * ones[:, None], axis=-1 - ) - grid_ = B.reshape(grid, _GRID_SIZE**2, 2) - - X = grid_ - - eigenfunctions = product.get_eigenfunctions(_TRUNC_LEVEL**2) - - Phi_X = eigenfunctions(X) # [GS**2, M] - Phi_X2 = eigenfunctions(X) - - weights = from_numpy(X, np.random.randn(eigenfunctions.num_levels)) - chained_weights = chain(weights, eigenfunctions.num_eigenfunctions_per_level) - weights = B.expand_dims(weights, -1) - actual = B.to_numpy(eigenfunctions.weighted_outerproduct(weights, X, X)) - - expected = B.einsum("ni,mi,i->nm", Phi_X, Phi_X2, chained_weights) - np.testing.assert_array_almost_equal(actual, expected) - - -def test_circle_product_kernel(): - product = ProductDiscreteSpectrumSpace( - Circle(), Circle(), num_levels=_TRUNC_LEVEL**2 - ) - - grid = B.linspace(0, 2 * B.pi, _GRID_SIZE) - ones = B.ones(_GRID_SIZE) - grid = B.stack( - grid[:, None] * ones[None, :], grid[None, :] * ones[:, None], axis=-1 - ) - grid_ = B.reshape(grid, _GRID_SIZE**2, 2) - - for ls in [0.1, 0.5, 1.0, 2.0, 5.0]: - kernel = MaternKarhunenLoeveKernel(product, _TRUNC_LEVEL**2) - kernel_single = MaternKarhunenLoeveKernel(Circle(), _TRUNC_LEVEL) - - params = kernel.init_params() - params["nu"] = from_numpy(grid_, [np.inf]) - params["lengthscale"] = from_numpy(grid, [ls]) - - params_single = kernel_single.init_params() - params_single["nu"] = from_numpy(grid_, [np.inf]) - params_single["lengthscale"] = from_numpy(grid, [ls]) - - k_xx = kernel.K(params, grid_, grid_) - k_xx = k_xx.reshape(_GRID_SIZE, _GRID_SIZE, _GRID_SIZE, _GRID_SIZE) - - k_xx_single_1 = kernel_single.K( - params_single, grid_[..., :1], grid_[..., :1] - ).reshape(_GRID_SIZE, _GRID_SIZE, _GRID_SIZE, _GRID_SIZE) - - k_xx_single_2 = kernel_single.K( - params_single, grid_[..., 1:], grid_[..., 1:] - ).reshape(_GRID_SIZE, _GRID_SIZE, _GRID_SIZE, _GRID_SIZE) - - k_xx_product = k_xx_single_1 * k_xx_single_2 - - np.testing.assert_allclose( - B.to_numpy(k_xx), B.to_numpy(k_xx_product), atol=1e-08, rtol=1e-05 - ) - - -def test_product_space_circle_su(): - circle = Circle() - su = SpecialUnitary(2) - - product = ProductDiscreteSpectrumSpace( - circle, - su, - num_levels=400, - num_levels_per_space=20, - ) - - key = B.create_random_state(np.float32) - key, xs_circle = circle.random(key, 1000) - key, xs_su = su.random(key, 1000) - - xs = make_product([xs_circle, xs_su]) - - kernel = MaternKarhunenLoeveKernel(product, 400) - kernel_single_circle = MaternKarhunenLoeveKernel(circle, 20) - kernel_single_su = MaternKarhunenLoeveKernel(su, 20) - - for ls in [0.1, 0.5, 1.0, 2.0, 5.0]: - - params = kernel.init_params() - params["nu"] = np.r_[np.inf] - params["lengthscale"] = np.r_[ls] - - k_xx = kernel.K(params, xs, xs[:1]) # [N, 1] - - k_xx_circle = kernel_single_circle.K(params, xs_circle, xs_circle[:1]) # [N, 1] - - k_xx_su = kernel_single_su.K(params, xs_su, xs_su[:1]) # [N, 1] - - k_xx_product = k_xx_circle * k_xx_su - - np.testing.assert_allclose(k_xx, k_xx_product, atol=1e-08, rtol=1e-05) - - -def test_product_space_circle_su_and_product_kernel(): - circle = Circle() - su = SpecialUnitary(2) - - product = ProductDiscreteSpectrumSpace( - circle, - su, - num_levels=400, - num_levels_per_space=20, - ) - - key = B.create_random_state(np.float32) - key, xs_circle = circle.random(key, 1000) - key, xs_su = su.random(key, 1000) - - xs = make_product([xs_circle, xs_su]) - - kernel = MaternKarhunenLoeveKernel(product, 400) - - kernel_single_circle = MaternKarhunenLoeveKernel(circle, 20) - kernel_single_su = MaternKarhunenLoeveKernel(su, 20) - - product_kernel = ProductGeometricKernel(kernel_single_circle, kernel_single_su) - - for ls in [0.1, 0.5, 1.0, 2.0, 5.0]: - - params = kernel.init_params() - params["nu"] = np.r_[np.inf] - params["lengthscale"] = np.r_[ls] - product_params = { - "nu": np.r_[np.inf, np.inf], - "lengthscale": np.r_[ls, ls], - } - - k_xx = kernel.K(params, xs, xs[:1]) # [N, 1] - k_xx_product = product_kernel.K(product_params, xs, xs[:1]) # [N, 1] - - np.testing.assert_allclose(k_xx, k_xx_product, atol=1e-08, rtol=1e-05) - - -def test_number_of_individual_eigenfunctions(): - circle = Circle() - sphere = Hypersphere(3) - - product = ProductDiscreteSpectrumSpace( - circle, - sphere, - num_levels=5, - num_levels_per_space=20, - ) - - eigf = product.get_eigenfunctions(5) - - key = B.create_random_state(np.float32) - N = 10 - key, xs_circle = circle.random(key, N) - key, xs_sph = sphere.random(key, N) - - xs = make_product([xs_circle, xs_sph]) - - assert eigf(xs).shape == (N, eigf.num_eigenfunctions) diff --git a/tests/kernels/test_product_kernel.py b/tests/kernels/test_product_kernel.py new file mode 100644 index 00000000..da21a402 --- /dev/null +++ b/tests/kernels/test_product_kernel.py @@ -0,0 +1,46 @@ +import numpy as np +import pytest + +from geometric_kernels.kernels import MaternGeometricKernel, ProductGeometricKernel +from geometric_kernels.spaces import Circle, SpecialUnitary +from geometric_kernels.utils.product import make_product + +from ..helper import check_function_with_backend + + +@pytest.mark.parametrize( + "factor1, factor2", [(Circle(), Circle()), (Circle(), SpecialUnitary(2))], ids=str +) +@pytest.mark.parametrize("nu, lengthscale", [(1 / 2, 2.0), (5 / 2, 1.0), (np.inf, 0.1)]) +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_kernel_is_product_of_heat_kernels(factor1, factor2, nu, lengthscale, backend): + key = np.random.RandomState(0) + key, xs_factor1 = factor1.random(key, 10) + key, xs_factor2 = factor2.random(key, 10) + + kernel_factor1 = MaternGeometricKernel(factor1) + kernel_factor2 = MaternGeometricKernel(factor2) + product_kernel = ProductGeometricKernel(kernel_factor1, kernel_factor2) + + def K_diff(nu, lengthscale, xs_factor1, xs_factor2): + params = {"nu": nu, "lengthscale": lengthscale} + + xs_product = make_product([xs_factor1, xs_factor2]) + + K_product = product_kernel.K(params, xs_product, xs_product) + K_factor1 = kernel_factor1.K(params, xs_factor1, xs_factor1) + K_factor2 = kernel_factor2.K(params, xs_factor2, xs_factor2) + + return K_product - K_factor1 * K_factor2 + + # Check that ProductGeometricKernel without ARD coincides with the product + # of the respective factor kernels. + check_function_with_backend( + backend, + np.zeros((xs_factor1.shape[0], xs_factor2.shape[0])), + K_diff, + np.array([nu]), + np.array([lengthscale]), + xs_factor1, + xs_factor2, + ) diff --git a/tests/kernels/test_sphere_heat_kernel.py b/tests/kernels/test_sphere_heat_kernel.py deleted file mode 100644 index 49890a0e..00000000 --- a/tests/kernels/test_sphere_heat_kernel.py +++ /dev/null @@ -1,58 +0,0 @@ -import lab as B -import numpy as np -import torch - -import geometric_kernels.torch # noqa -from geometric_kernels.kernels import MaternKarhunenLoeveKernel -from geometric_kernels.spaces.hypersphere import Hypersphere -from geometric_kernels.utils.manifold_utils import manifold_laplacian - -_TRUNCATION_LEVEL = 10 -_NU = 2.5 - - -def test_sphere_heat_kernel(): - # Parameters - grid_size = 4 - nb_samples = 10 - dimension = 3 - - # Create manifold - hypersphere = Hypersphere(dim=dimension) - - # Generate samples - ts = torch.linspace(0.1, 1, grid_size, requires_grad=True) - xs = torch.tensor( - np.array(hypersphere.random_point(nb_samples)), requires_grad=True - ) - ys = xs - - # Define kernel - kernel = MaternKarhunenLoeveKernel(hypersphere, _TRUNCATION_LEVEL, normalize=False) - params = kernel.init_params() - params["nu"] = torch.tensor([torch.inf]) - - # Define heat kernel function - def heat_kernel(t, x, y): - params["lengthscale"] = B.reshape(B.sqrt(2 * t), 1) - return kernel.K(params, x, y) - - for t in ts: - for x in xs: - for y in ys: - # Compute the derivative of the kernel function wrt t - dfdt, _, _ = torch.autograd.grad( - heat_kernel(t, x[None], y[None]), (t, x, y) - ) - # Compute the Laplacian of the kernel on the manifold - egrad = lambda u: torch.autograd.grad( # noqa - heat_kernel(t, u[None], y[None]), (t, u, y) - )[ - 1 - ] # noqa - fx = lambda u: heat_kernel(t, u[None], y[None]) # noqa - ehess = lambda u, h: torch.autograd.functional.hvp(fx, u, h)[1] # noqa - lapf = manifold_laplacian(x, hypersphere, egrad, ehess) - - # Check that they match - assert np.isclose(dfdt.detach().numpy(), lapf, atol=1.0e-3) diff --git a/tests/sampling/test_samplers.py b/tests/sampling/test_samplers.py new file mode 100644 index 00000000..557654d0 --- /dev/null +++ b/tests/sampling/test_samplers.py @@ -0,0 +1,34 @@ +import lab as B +import numpy as np +import pytest + +from geometric_kernels.sampling import sampler + +from ..feature_maps.test_feature_maps import feature_map_and_friends # noqa: F401 +from ..helper import check_function_with_backend, create_random_state + +_NUM_SAMPLES = 2 + + +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_output_shape_and_backend(backend, feature_map_and_friends): + feature_map, kernel, space = feature_map_and_friends + + params = kernel.init_params() + sample_paths = sampler(feature_map, s=_NUM_SAMPLES) + + key = np.random.RandomState(0) + key, X = space.random(key, 50) + + def sample(params, X): + return sample_paths(X, params, key=create_random_state(backend))[1] + + # Check that sample_paths runs and the output is a tensor of the right backend and shape. + check_function_with_backend( + backend, + (X.shape[0], _NUM_SAMPLES), + sample, + params, + X, + compare_to_result=lambda res, f_out: B.shape(f_out) == res, + ) diff --git a/tests/spaces/test_basics.py b/tests/spaces/test_basics.py new file mode 100644 index 00000000..a0e300ee --- /dev/null +++ b/tests/spaces/test_basics.py @@ -0,0 +1,52 @@ +import inspect + +import pytest + +import geometric_kernels.spaces + +from ..helper import ( + compact_matrix_lie_groups, + discrete_spectrum_spaces, + noncompact_symmetric_spaces, + product_discrete_spectrum_spaces, + spaces, +) + + +@pytest.mark.parametrize( + "fun, cls", + [ + (compact_matrix_lie_groups, geometric_kernels.spaces.CompactMatrixLieGroup), + ( + product_discrete_spectrum_spaces, + geometric_kernels.spaces.ProductDiscreteSpectrumSpace, + ), + (discrete_spectrum_spaces, geometric_kernels.spaces.DiscreteSpectrumSpace), + ( + noncompact_symmetric_spaces, + geometric_kernels.spaces.NoncompactSymmetricSpace, + ), + (spaces, geometric_kernels.spaces.Space), + ], +) +def test_all_discrete_spectrum_spaces_covered(fun, cls): + spaces = fun() + + # all classes in the geometric_kernels.spaces module + classes = [ + (cls_name, cls_obj) + for cls_name, cls_obj in inspect.getmembers(geometric_kernels.spaces) + if inspect.isclass(cls_obj) + ] + for cls_name, cls_obj in classes: + if issubclass(cls_obj, cls) and not inspect.isabstract(cls_obj): + for space in spaces: + if isinstance(space, cls_obj): + break + else: + # complain if discrete_spectrum_spaces() does not contain an + # instance of a non-abstract subclass of DiscreteSpectrumSpace + # from the geometric_kernels.spaces module. + assert ( + False + ), f"An instance of the class `{cls_name}` is missing from the list returned by the function `{fun.__name__}`" diff --git a/tests/spaces/test_circle.py b/tests/spaces/test_circle.py index 426ff2a5..c7e59d7a 100644 --- a/tests/spaces/test_circle.py +++ b/tests/spaces/test_circle.py @@ -1,175 +1,54 @@ import lab as B import numpy as np import pytest -import tensorflow as tf -import torch -from plum import Tuple -from geometric_kernels.kernels import MaternKarhunenLoeveKernel -from geometric_kernels.lab_extras import from_numpy -from geometric_kernels.spaces.circle import Circle, SinCosEigenfunctions -from geometric_kernels.spaces.eigenfunctions import EigenfunctionsWithAdditionTheorem -from geometric_kernels.utils.utils import chain +from geometric_kernels.kernels import MaternGeometricKernel +from geometric_kernels.spaces.circle import Circle +from geometric_kernels.utils.kernel_formulas import ( + euclidean_matern_12_kernel, + euclidean_matern_32_kernel, + euclidean_matern_52_kernel, + euclidean_rbf_kernel, +) +from ..helper import check_function_with_backend -class Consts: - seed = 42 - num_data = 7 - num_data2 = 5 - num_eigenfunctions = 11 - num_levels = 6 - -def to_typed_tensor(value, backend): - if backend == "tensorflow": - return tf.convert_to_tensor(value) - elif backend == "torch": - return torch.tensor(value) - elif backend == "numpy": - return value - else: - raise ValueError("Unknown backend: {}".format(backend)) - - -@pytest.fixture(name="inputs", params=["tensorflow", "torch", "numpy"]) -def _inputs_fixure(request) -> Tuple[B.Numeric]: - np.random.seed(Consts.seed) - value = np.random.uniform(0, 2 * np.pi, size=(Consts.num_data, 1)) - value2 = np.random.uniform(0, 2 * np.pi, size=(Consts.num_data2, 1)) - return to_typed_tensor(value, request.param), to_typed_tensor(value2, request.param) - - -@pytest.fixture(name="eigenfunctions") -def _eigenfunctions_fixture(): - eigenfunctions = SinCosEigenfunctions(Consts.num_levels) - return eigenfunctions - - -def test_call_eigenfunctions( - inputs: Tuple[B.Numeric, B.Numeric], - eigenfunctions: EigenfunctionsWithAdditionTheorem, -): - inputs, _ = inputs - output = B.to_numpy(eigenfunctions(inputs)) - assert output.shape == (Consts.num_data, eigenfunctions.num_eigenfunctions) - - -def test_eigenfunctions_shape(eigenfunctions: EigenfunctionsWithAdditionTheorem): - num_eigenfunctions_manual = np.sum(eigenfunctions.num_eigenfunctions_per_level) - assert num_eigenfunctions_manual == eigenfunctions.num_eigenfunctions - assert len(eigenfunctions.num_eigenfunctions_per_level) == eigenfunctions.num_levels - - -def test_orthonormality(eigenfunctions: EigenfunctionsWithAdditionTheorem): - theta = np.linspace(-np.pi, np.pi, 5_000).reshape(-1, 1) # [N, 1] - phi = B.to_numpy(eigenfunctions(theta)) - phiT_phi = (phi.T @ phi) * 2 * np.pi / phi.shape[0] - circumference_circle = 2 * np.pi - inner_prod = phiT_phi / circumference_circle - np.testing.assert_array_almost_equal( - inner_prod, np.eye(inner_prod.shape[0]), decimal=3 - ) - - -def test_weighted_outerproduct_with_addition_theorem( - inputs, eigenfunctions: EigenfunctionsWithAdditionTheorem -): - """ - Eigenfunction will use addition theorem to compute outerproduct. We compare against the - naive implementation. - """ - inputs, inputs2 = inputs - weights_per_level = from_numpy(inputs, np.random.randn(eigenfunctions.num_levels)) - weights = B.reshape(weights_per_level, -1, 1) - chained_weights = chain( - weights_per_level, eigenfunctions.num_eigenfunctions_per_level - ) - actual = B.to_numpy(eigenfunctions.weighted_outerproduct(weights, inputs, inputs2)) - - Phi_X = eigenfunctions(inputs) - Phi_X2 = eigenfunctions(inputs2) - expected = B.einsum("ni,ki,i->nk", Phi_X, Phi_X2, chained_weights) - np.testing.assert_array_almost_equal(actual, expected) - - -def test_weighted_outerproduct_with_addition_theorem_same_input( - inputs, eigenfunctions: EigenfunctionsWithAdditionTheorem -): - """ - Eigenfunction will use addition theorem to compute outerproduct. We compare against the - naive implementation. - """ - inputs, _ = inputs - weights_per_level = from_numpy(inputs, np.random.randn(eigenfunctions.num_levels)) - weights = B.reshape(weights_per_level, -1, 1) - first = B.to_numpy(eigenfunctions.weighted_outerproduct(weights, inputs, inputs)) - second = B.to_numpy(eigenfunctions.weighted_outerproduct(weights, inputs, None)) - np.testing.assert_array_almost_equal(first, second) - - -def test_weighted_outerproduct_diag_with_addition_theorem( - inputs, eigenfunctions: EigenfunctionsWithAdditionTheorem -): - """ - Eigenfunction will use addition theorem to compute outerproduct. We compare against the - naive implementation. - """ - inputs, _ = inputs - weights_per_level = from_numpy(inputs, np.random.randn(eigenfunctions.num_levels)) - chained_weights = chain( - weights_per_level, eigenfunctions.num_eigenfunctions_per_level - ) - weights = B.reshape(weights_per_level, -1, 1) - actual = eigenfunctions.weighted_outerproduct_diag(weights, inputs) - - Phi_X = eigenfunctions(inputs) - expected = B.einsum("ni,i->n", Phi_X**2, chained_weights) - np.testing.assert_array_almost_equal(B.to_numpy(actual), B.to_numpy(expected)) - - -def analytic_kernel(nu: float, r: B.Numeric) -> B.Numeric: - """ - Analytic implementations of matern-family kernels. - - :param nu: selects the matern - :param r: distance, shape [...] - :return: k(r), shape [...] - """ - r = B.abs(r) +@pytest.mark.parametrize( + "nu, atol", [(0.5, 1e-1), (1.5, 1e-3), (2.5, 1e-3), (np.inf, 1e-5)] +) +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_equivalence_kernel(nu, atol, backend): if nu == 0.5: - return B.exp(-r) + analytic_kernel = euclidean_matern_12_kernel elif nu == 1.5: - sqrt3 = np.sqrt(3.0) - return (1.0 + sqrt3 * r) * B.exp(-sqrt3 * r) + analytic_kernel = euclidean_matern_32_kernel elif nu == 2.5: - sqrt5 = np.sqrt(5.0) - return (1.0 + sqrt5 * r + 5.0 / 3.0 * (r**2)) * B.exp(-sqrt5 * r) + analytic_kernel = euclidean_matern_52_kernel elif nu == np.inf: - return B.exp(-0.5 * r**2) - else: - raise NotImplementedError + analytic_kernel = euclidean_rbf_kernel + inputs = np.random.uniform(0, 2 * np.pi, size=(5, 1)) + inputs2 = np.random.uniform(0, 2 * np.pi, size=(3, 1)) -@pytest.mark.parametrize("nu, decimal", [(0.5, 1), (1.5, 3), (2.5, 5), (np.inf, 6)]) -def test_equivalence_kernel(nu, decimal, inputs): - inputs, inputs2 = inputs - # Spectral kernel - circle = Circle() - kernel = MaternKarhunenLoeveKernel(circle, num_levels=101) - params = kernel.init_params() - params["nu"] = from_numpy(inputs, np.r_[nu]) - params["lengthscale"] = from_numpy(inputs, np.r_[1.0]) - - K_actual = B.to_numpy(kernel.K(params, inputs, inputs2)) - - # Kernel by summing over all distances + # Compute kernel using periodic summation geodesic = inputs[:, None, :] - inputs2[None, :, :] # [N, N2, 1] all_distances = ( geodesic + np.array([i * 2 * np.pi for i in range(-10, 10)])[None, None, :] ) - K_expected = B.to_numpy(B.sum(analytic_kernel(nu, all_distances), axis=2)) - - # test equivalence - np.testing.assert_array_almost_equal( - K_expected / K_expected[0, 0], K_actual / K_actual[0, 0], decimal=decimal + all_distances = B.abs(all_distances) + result = B.to_numpy(B.sum(analytic_kernel(all_distances), axis=2)) + + kernel = MaternGeometricKernel(Circle()) + + # Check that MaternGeometricKernel on Circle() coincides with the + # periodic summation of the respective Euclidean Matérn kernel. + check_function_with_backend( + backend, + result, + kernel.K, + {"nu": np.array([nu]), "lengthscale": np.array([1.0])}, + inputs, + inputs2, + atol=atol, ) diff --git a/tests/spaces/test_eigenfunctions.py b/tests/spaces/test_eigenfunctions.py new file mode 100644 index 00000000..c6a98b93 --- /dev/null +++ b/tests/spaces/test_eigenfunctions.py @@ -0,0 +1,202 @@ +import jax +import lab as B +import numpy as np +import pytest +from opt_einsum import contract as einsum + +from geometric_kernels.kernels.matern_kernel import default_num +from geometric_kernels.spaces import CompactMatrixLieGroup, Hypersphere, Mesh +from geometric_kernels.utils.utils import chain + +from ..helper import check_function_with_backend, discrete_spectrum_spaces + +jax.config.update("jax_enable_x64", True) # enable float64 in JAX + + +@pytest.fixture( + params=discrete_spectrum_spaces(), + ids=str, +) +def inputs(request): + """ + Returns a tuple (space, eigenfunctions, X, X2, weights) where: + - space = request.param, + - eigenfunctions = space.get_eigenfunctions(num_levels), with reasonable num_levels + - X is a random sample of random size from the space, + - X2 is another random sample of random size from the space, + - weights is an array of positive numbers of shape (eigenfunctions.num_levels, 1). + """ + space = request.param + num_levels = default_num(space) + if isinstance(space, Hypersphere): + # For Hypersphere, the maximal number of levels with eigenfunction + # evaluation support is stored in the num_computed_levels field. We + # do not use more levels to be able to test eigenfunction evaluation. + num_levels = min( + 10, num_levels, space.get_eigenfunctions(num_levels).num_computed_levels + ) + elif isinstance(space, Mesh): + # We limit the number of levels to 50 to avoid excessive computation + # and increased numerical error for higher order eigenfunctions. + num_levels = min(50, num_levels) + eigenfunctions = space.get_eigenfunctions(num_levels) + + key = np.random.RandomState(0) + N, N2 = key.randint(low=1, high=100 + 1, size=2) + key, X = space.random(key, N) + key, X2 = space.random(key, N2) + + # These weights are used for testing the weighted outerproduct, they + # should be positive. + weights = np.random.rand(eigenfunctions.num_levels, 1) ** 2 + 1e-5 + + return space, eigenfunctions, X, X2, weights + + +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_call_eigenfunctions(inputs, backend): + space, eigenfunctions, X, _, _ = inputs + + if isinstance(space, CompactMatrixLieGroup): + pytest.skip( + "CompactMatrixLieGroup subclasses do not currently support eigenfunction evaluation" + ) + + # Check that the eigenfunctions can be called, returning the right type and shape. + check_function_with_backend( + backend, + (X.shape[0], eigenfunctions.num_eigenfunctions), + eigenfunctions, + X, + compare_to_result=lambda res, f_out: f_out.shape == res, + ) + + +def test_numbers_of_eigenfunctions(inputs): + _, eigenfunctions, _, _, _ = inputs + num_levels = eigenfunctions.num_levels + # Check that the length of the `num_eigenfunctions_per_level` list is correct. + assert len(eigenfunctions.num_eigenfunctions_per_level) == num_levels + # Check that the first eigenspace is 1-dimensional. + assert eigenfunctions.num_eigenfunctions_per_level[0] == 1 + + # Check that dimensions of eigenspaces are always positive. + for i in range(num_levels): + assert eigenfunctions.num_eigenfunctions_per_level[i] > 0 + + num_eigenfunctions_manual = sum(eigenfunctions.num_eigenfunctions_per_level) + # Check that `num_eigenfunctions_per_level` sum up to the total number of + # eigenfunctions. + assert num_eigenfunctions_manual == eigenfunctions.num_eigenfunctions + + +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_orthonormality(inputs, backend): + space, eigenfunctions, _, _, _ = inputs + + if isinstance(space, CompactMatrixLieGroup): + pytest.skip( + "CompactMatrixLieGroup subclasses do not currently support eigenfunction evaluation" + ) + + key = np.random.RandomState(0) + key, xs = space.random(key, 10000) + + # Check that the eigenfunctions are orthonormal by comparing a Monte Carlo + # approximation of the inner product with the identity matrix. + check_function_with_backend( + backend, + np.eye(eigenfunctions.num_eigenfunctions), + lambda xs: B.matmul(B.T(eigenfunctions(xs)), eigenfunctions(xs)) / xs.shape[0], + xs, + atol=0.4, # very loose, but helps make sure the diagonal is close to 1 while the rest is close to 0 + ) + + +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_weighted_outerproduct_with_addition_theorem(inputs, backend): + space, eigenfunctions, X, X2, weights = inputs + + if isinstance(space, CompactMatrixLieGroup): + pytest.skip( + "CompactMatrixLieGroup subclasses do not currently support eigenfunction evaluation" + ) + + chained_weights = chain( + weights.squeeze(), eigenfunctions.num_eigenfunctions_per_level + ) + + Phi_X = eigenfunctions(X) + Phi_X2 = eigenfunctions(X2) + result = einsum("ni,ki,i->nk", Phi_X, Phi_X2, chained_weights) + + # Check that `weighted_outerproduct`, which is based on the addition theorem, + # returns the same result as the direct computation involving individual + # eigenfunctions. + check_function_with_backend( + backend, result, eigenfunctions.weighted_outerproduct, weights, X, X2, atol=1e-2 + ) + + +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_weighted_outerproduct_with_addition_theorem_one_input(inputs, backend): + _, eigenfunctions, X, _, weights = inputs + + result = eigenfunctions.weighted_outerproduct(weights, X, X) + + # Check that `weighted_outerproduct`, when given only X (but not X2), + # uses X2=X. + check_function_with_backend( + backend, + result, + eigenfunctions.weighted_outerproduct, + weights, + X, + atol=1e-2, + ) + + +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_weighted_outerproduct_diag(inputs, backend): + _, eigenfunctions, X, _, weights = inputs + + result = np.diag(eigenfunctions.weighted_outerproduct(weights, X, X)) + + # Check that `weighted_outerproduct_diag` returns the same result as the + # diagonal of the full `weighted_outerproduct`. + check_function_with_backend( + backend, + result, + eigenfunctions.weighted_outerproduct_diag, + weights, + X, + ) + + +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_weighted_outerproduct_against_phi_product(inputs, backend): + _, eigenfunctions, X, X2, weights = inputs + + sum_phi_phi_for_level = eigenfunctions.phi_product(X, X2) + + result = np.einsum("id,...nki->...nk", weights, sum_phi_phi_for_level) + + # Check that `weighted_outerproduct` returns the weighted sum of `phi_product`. + check_function_with_backend( + backend, result, eigenfunctions.weighted_outerproduct, weights, X, X2 + ) + + +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_weighted_outerproduct_diag_against_phi_product(inputs, backend): + _, eigenfunctions, X, _, weights = inputs + + phi_product_diag = eigenfunctions.phi_product_diag(X) + + result = np.einsum("id,ni->n", weights, phi_product_diag) # [N,] + + # Check that `weighted_outerproduct_diag` returns the weighted sum + # of `phi_product_diag`. + check_function_with_backend( + backend, result, eigenfunctions.weighted_outerproduct_diag, weights, X + ) diff --git a/tests/spaces/test_eigenvalues.py b/tests/spaces/test_eigenvalues.py new file mode 100644 index 00000000..6b04fb8b --- /dev/null +++ b/tests/spaces/test_eigenvalues.py @@ -0,0 +1,54 @@ +""" +Note: We don't use `check_function_with_backend` throughout this module because +eigenvalues are always represented by numpy arrays, regardless of the backend +used for other routines. +""" + +import numpy as np +import pytest + +from geometric_kernels.kernels.matern_kernel import default_num + +from ..helper import discrete_spectrum_spaces + + +@pytest.fixture( + params=discrete_spectrum_spaces(), + ids=str, +) +def inputs(request): + """ + Returns a tuple (space, num_levels, eigenvalues) where: + - space = request.param, + - num_levels is the default number of levels for the `space`, if it does not + exceed 100, and 100 otherwise, + - eigenvalues = space.get_eigenvalues(num_levels), + - eps, a small number, a technicality for using `assert_array_less`. + """ + space = request.param + num_levels = min(default_num(space), 100) + eigenvalues = space.get_eigenvalues(num_levels) + eps = 1e-5 + + return space, num_levels, eigenvalues, eps + + +def test_shape(inputs): + _, num_levels, eigenvalues, _ = inputs + + # Check that the eigenvalues have appropriate shape. + assert eigenvalues.shape == (num_levels, 1) + + +def test_positive(inputs): + _, _, eigenvalues, eps = inputs + + # Check that the eigenvalues are nonnegative. + np.testing.assert_array_less(np.zeros_like(eigenvalues), eigenvalues + eps) + + +def test_ordered(inputs): + _, _, eigenvalues, eps = inputs + + # Check that the eigenvalues are sorted in ascending order. + np.testing.assert_array_less(eigenvalues[:-1], eigenvalues[1:] + eps) diff --git a/tests/spaces/test_graph.py b/tests/spaces/test_graph.py index fc94155c..a31fa42a 100644 --- a/tests/spaces/test_graph.py +++ b/tests/spaces/test_graph.py @@ -1,220 +1,122 @@ import warnings -import jax import lab as B import numpy as np -import scipy.sparse as sp -import tensorflow as tf -import torch +import pytest from geometric_kernels.jax import * # noqa -from geometric_kernels.kernels import MaternKarhunenLoeveKernel +from geometric_kernels.kernels import MaternGeometricKernel from geometric_kernels.spaces import Graph from geometric_kernels.tensorflow import * # noqa from geometric_kernels.torch import * # noqa -warnings.filterwarnings("ignore", category=RuntimeWarning, module="scipy") - -A = np.array( - [ - [0, 1, 0, 0, 0, 0, 0], - [1, 0, 1, 1, 1, 0, 0], - [0, 1, 0, 0, 0, 1, 0], - [0, 1, 0, 0, 1, 0, 0], - [0, 1, 0, 1, 0, 0, 0], - [0, 0, 1, 0, 0, 0, 0], - [0, 0, 0, 0, 0, 0, 0], - ] -).astype(np.float64) - -L = np.array( - [ - [1, -1, 0, 0, 0, 0, 0], - [-1, 4, -1, -1, -1, 0, 0], - [0, -1, 2, 0, 0, -1, 0], - [0, -1, 0, 2, -1, 0, 0], - [0, -1, 0, -1, 2, 0, 0], - [0, 0, -1, 0, 0, 1, 0], - [0, 0, 0, 0, 0, 0, 0], - ] -).astype(np.float64) - - -def run_tests_with_adj(A, L, tol=1e-7, tol_m=1e-4): - ############################################## - # Inits - - n = A.shape[0] - L = B.cast(B.dtype(A), L) - graph = Graph(A) - - normed_graph = Graph(A, normalize_laplacian=True) - - ############################################## - # Laplacian computation - - comparison = graph._laplacian == L - - if sp.issparse(comparison): - comparison = comparison.toarray() - - if isinstance(comparison, np.matrix): # bug with lab? - assert comparison.all(), "Laplacian does not match." - else: - assert B.all(comparison), "Laplacian does not match." - - normed_l = normed_graph._laplacian - if sp.issparse(normed_l): - normed_l = normed_l.toarray() - - assert ( - B.max(B.abs(B.diag(normed_l) - 1)[:-1]) < tol_m - and B.abs(B.diag(normed_l)[-1] - 0) < tol_m - ) +from ..data import ( + TEST_GRAPH_ADJACENCY, + TEST_GRAPH_LAPLACIAN, + TEST_GRAPH_LAPLACIAN_NORMALIZED, +) +from ..helper import check_function_with_backend, np_to_backend - ############################################## - # Eigendecomposition checks +warnings.filterwarnings("ignore", category=RuntimeWarning, module="scipy") - evecs = graph.get_eigenvectors(n) - evals = graph.get_eigenvalues(n) +A = TEST_GRAPH_ADJACENCY +L = TEST_GRAPH_LAPLACIAN - evals_np, evecs_np = np.linalg.eigh(L) - evecs_np *= np.sqrt(graph.num_vertices) - # check vals - np.testing.assert_allclose(evals[:, 0], evals_np, atol=tol, rtol=tol) +@pytest.mark.parametrize("normalized", [True, False]) +@pytest.mark.parametrize( + "backend", ["numpy", "tensorflow", "torch", "jax", "scipy_sparse"] +) +def test_laplacian(normalized, backend): - # check vecs - np.testing.assert_allclose( - np.abs(evecs)[:, 2:], np.abs(evecs_np)[:, 2:], atol=tol, rtol=tol + # Check that the Laplacian is computed correctly. + check_function_with_backend( + backend, + TEST_GRAPH_LAPLACIAN if not normalized else TEST_GRAPH_LAPLACIAN_NORMALIZED, + lambda adj: Graph(adj, normalize_laplacian=normalized)._laplacian, + TEST_GRAPH_ADJACENCY, ) - try: - np.testing.assert_allclose( - np.abs(evecs)[:, :2], np.abs(evecs_np)[:, :2], atol=tol, rtol=tol - ) - except AssertionError: - np.testing.assert_allclose( - np.abs(evecs)[:, [1, 0]], np.abs(evecs_np)[:, :2], atol=tol, rtol=tol - ) - - normed_evals = normed_graph.get_eigenvalues(n) - assert (B.min(normed_evals) >= 0) and ( - B.max(normed_evals) <= 2 - ) # well known inequality - ############################################## - # Kernel init checks - - K_cons = MaternKarhunenLoeveKernel(graph, n, normalize=False) - params = K_cons.init_params() - - idx = B.cast(B.dtype(A), np.arange(n)[:, None]) - - nu = B.cast(B.dtype(A), np.array([1.0])) - lscale = B.cast(B.dtype(A), np.array([1.0])) - - K_normed_cons = MaternKarhunenLoeveKernel(normed_graph, n, normalize=False) - normed_params = K_normed_cons.init_params() - - ############################################## - # Matern 1 check +@pytest.mark.parametrize( + "L", [TEST_GRAPH_ADJACENCY.shape[0], TEST_GRAPH_ADJACENCY.shape[0] // 2] +) +@pytest.mark.parametrize("normalized", [True, False]) +@pytest.mark.parametrize( + "backend", ["numpy", "tensorflow", "torch", "jax", "scipy_sparse"] +) +def test_eigendecomposition(L, normalized, backend): + laplacian = np_to_backend( + TEST_GRAPH_LAPLACIAN if not normalized else TEST_GRAPH_LAPLACIAN_NORMALIZED, + backend, + ) - params["nu"], params["lengthscale"] = nu, lscale - Kg = K_cons.K(params, idx) - K1 = evecs_np @ np.diag(np.power(evals_np + 2, -1)) @ evecs_np.T - np.testing.assert_allclose(Kg, K1, atol=tol, rtol=tol) + def eigendiff(adj): + graph = Graph(adj, normalize_laplacian=normalized) - normed_params["nu"], normed_params["lengthscale"] = nu, lscale - K_normed_cons.K(normed_params, idx) + eigenvalue_mat = B.diag_construct(graph.get_eigenvalues(L)[:, 0]) + eigenvectors = graph.get_eigenvectors(L) + # If the backend is scipy_sparse, convert eigenvalues/eigenvectors, + # which are always supposed to be dense arrays, to sparse arrays. + if backend == "scipy_sparse": + import scipy.sparse as sp - ############################################## - # Matern 2 check + eigenvalue_mat = sp.csr_array(eigenvalue_mat) + eigenvectors = sp.csr_array(eigenvectors) - nu = B.cast(B.dtype(A), np.array([2.0])) - params["nu"], params["lengthscale"] = nu, lscale - Kg = K_cons.K(params, idx) - K2 = evecs_np @ np.diag(np.power(evals_np + 4, -2)) @ evecs_np.T - np.testing.assert_allclose(Kg, K2, atol=tol, rtol=tol) + laplace_x_eigvecs = laplacian @ eigenvectors + eigvals_x_eigvecs = eigenvectors @ eigenvalue_mat + return laplace_x_eigvecs - eigvals_x_eigvecs - ############################################## - # RBF check + check_function_with_backend( + backend, + np.zeros((TEST_GRAPH_ADJACENCY.shape[0], L)), + eigendiff, + TEST_GRAPH_ADJACENCY, + ) - nu = B.cast(B.dtype(A), np.array([np.inf])) - params["nu"], params["lengthscale"] = nu, lscale - Kg = K_cons.K(params, idx) - Ki = evecs_np @ np.diag(np.exp(-0.5 * evals_np)) @ evecs_np.T - np.testing.assert_allclose(Kg, Ki, atol=tol, rtol=tol) - ############################################## - # Fewer than n eigencomps check +@pytest.mark.parametrize("nu, lengthscale", [(1.0, 1.0), (2.0, 1.0), (np.inf, 1.0)]) +@pytest.mark.parametrize("sparse_adj", [True, False]) +@pytest.mark.parametrize("normalized", [True, False]) +@pytest.mark.parametrize( + "backend", ["numpy", "tensorflow", "torch", "jax"] +) # The kernels never take sparse parameters and never output sparse matrices, thus we don't test scipy_sparse. The fact that the adjacency matrix may be sparse is tested when sparse_adj is True. +def test_matern_kernels(nu, lengthscale, sparse_adj, normalized, backend): - m = 4 - evecs = graph.get_eigenvectors(m) - evals = graph.get_eigenvalues(m) - if isinstance(L, jax.numpy.ndarray): - evals_np, evecs_np = np.linalg.eigh(B.to_numpy(L)) - evals_np, evecs_np = evals_np[:m], evecs_np[:, :m] - else: - evals_np, evecs_np = sp.linalg.eigsh(B.to_numpy(L), m, sigma=1e-8) - evecs_np *= np.sqrt(graph.num_vertices) - - np.testing.assert_allclose(evals[:, 0], evals_np, atol=tol_m, rtol=tol_m) + laplacian = ( + TEST_GRAPH_LAPLACIAN if not normalized else TEST_GRAPH_LAPLACIAN_NORMALIZED + ) - try: - np.testing.assert_allclose( - np.abs(evecs)[:, :2], np.abs(evecs_np)[:, :2], atol=tol_m, rtol=tol_m + evals_np, evecs_np = np.linalg.eigh(laplacian) + evecs_np *= np.sqrt(laplacian.shape[0]) + + def evaluate_kernel(adj, nu, lengthscale): + dtype = B.dtype(adj) + if sparse_adj: + adj = np_to_backend(B.to_numpy(adj), "scipy_sparse") + graph = Graph(adj, normalize_laplacian=normalized) + kernel = MaternGeometricKernel(graph) + return kernel.K( + {"nu": nu, "lengthscale": lengthscale}, + B.range(dtype, adj.shape[0])[:, None], ) - except AssertionError: - np.testing.assert_allclose( - np.abs(evecs)[:, [1, 0]], np.abs(evecs_np)[:, :2], atol=tol_m, rtol=tol_m - ) - - K_cons = MaternKarhunenLoeveKernel(graph, m, normalize=False) - params = K_cons.init_params() - - nu = B.cast(B.dtype(A), np.array([np.inf])) - params["nu"], params["lengthscale"] = nu, lscale - Kg = K_cons.K(params, idx) - Ki = evecs_np @ np.diag(np.exp(-0.5 * evals_np)) @ evecs_np.T - np.testing.assert_allclose(Kg, Ki, atol=tol_m, rtol=tol_m) - - -def test_graphs_numpy(): - run_tests_with_adj(A, L) - - -def test_graphs_scipy_sparse(): - run_tests_with_adj(sp.csr_matrix(A), L) - - -def test_graphs_torch(): - run_tests_with_adj(torch.tensor(A), L) - -def test_graphs_tf(): - run_tests_with_adj(tf.Variable(A), L) - - -def test_graphs_jax(): - run_tests_with_adj(jax.numpy.array(A), L, 1e-4, 1e-4) - - -def test_graphs_torch_cuda(): - if torch.cuda.is_available(): - adj = torch.tensor(A).cuda() - - n = adj.shape[0] - graph = Graph(adj) - # normed_graph = Graph(adj, normalize_laplacian=True) # fails due to bug in lab - - K_cons = MaternKarhunenLoeveKernel(graph, n, normalize=False) - params = K_cons.init_params() - - params["nu"] = torch.nn.Parameter(torch.tensor([1.0]).cuda()) - params["lengthscale"] = torch.nn.Parameter(torch.tensor([1.0]).cuda()) - - idx = torch.arange(n)[:, None].cuda() - K_cons.K(params, idx) + if nu < np.inf: + K = ( + evecs_np + @ np.diag(np.power(evals_np + 2 * nu / lengthscale**2, -nu)) + @ evecs_np.T + ) else: - pass + K = evecs_np @ np.diag(np.exp(-(lengthscale**2) / 2 * evals_np)) @ evecs_np.T + K = K / np.mean(K.diagonal()) + + check_function_with_backend( + backend, + K, + evaluate_kernel, + TEST_GRAPH_ADJACENCY, + np.array([nu]), + np.array([lengthscale]), + ) diff --git a/tests/spaces/test_hyperbolic.py b/tests/spaces/test_hyperbolic.py index 79e2d0ee..9b51260f 100644 --- a/tests/spaces/test_hyperbolic.py +++ b/tests/spaces/test_hyperbolic.py @@ -1,28 +1,53 @@ +import lab as B import numpy as np - -from geometric_kernels.spaces.hyperbolic import Hyperbolic - - -def test_hyperboloid_distance(): - hyperboloid = Hyperbolic(dim=2) - N = 10 - - # Data points - base = np.r_[7.14142843e00, -5.00000000e00, -5.00000000e00] - point = np.r_[14.17744688, 10.0, 10.0] - geodesic = hyperboloid.metric.geodesic(initial_point=base, end_point=point) - x1 = geodesic(np.linspace(0.0, 1.0, N)) # (N, 3) - x2 = x1[-1, None] # (1, 3) - - our_dist_12 = hyperboloid.distance(x2, x1) - geomstats_dist_12 = hyperboloid.metric.dist(x2, x1) - - our_dist_11 = hyperboloid.distance(x1, x1) # (N, N) - geomstats_dist_11 = hyperboloid.metric.dist_pairwise(x1, n_jobs=1) # (N, N) - - our_dist_11_diag = hyperboloid.distance(x1, x1, diag=True) # (N, ) - geomstats_dist_11_diag = hyperboloid.metric.dist(x1, x1) # (N, ) - - assert np.allclose(our_dist_12, geomstats_dist_12) - assert np.allclose(our_dist_11, geomstats_dist_11) - assert np.allclose(our_dist_11_diag, geomstats_dist_11_diag) +import pytest + +from geometric_kernels.kernels import MaternGeometricKernel +from geometric_kernels.spaces import Hyperbolic +from geometric_kernels.utils.kernel_formulas import ( + hyperbolic_heat_kernel_even, + hyperbolic_heat_kernel_odd, +) + +from ..helper import check_function_with_backend, create_random_state + + +@pytest.mark.parametrize("dim", [2, 3, 5, 7]) +@pytest.mark.parametrize("lengthscale", [2.0]) +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_equivalence_kernel(dim, lengthscale, backend): + space = Hyperbolic(dim) + + key = np.random.RandomState(0) + key, X = space.random(key, 6) + X2 = X.copy() + + t = lengthscale * lengthscale / 2 + if dim % 2 == 1: + result = hyperbolic_heat_kernel_odd(dim, t, X, X2) + else: + result = hyperbolic_heat_kernel_even(dim, t, X, X2) + + kernel = MaternGeometricKernel(space, key=create_random_state(backend)) + + def compare_to_result(res, f_out): + return ( + np.linalg.norm(res - B.to_numpy(f_out)) + / np.sqrt(res.shape[0] * res.shape[1]) + < 1e-1 + ) + + # Check that MaternGeometricKernel on Hyperbolic(dim) with nu=inf coincides + # with the well-known analytic formula for the heat kernel on the hyperbolic + # space in odd dimensions and semi-analytic formula in even dimensions. + # We are checking the equivalence on average, computing the norm between + # the two covariance matrices. + check_function_with_backend( + backend, + result, + kernel.K, + {"nu": np.array([np.inf]), "lengthscale": np.array([lengthscale])}, + X, + X2, + compare_to_result=compare_to_result, + ) diff --git a/tests/spaces/test_hypercube_graph.py b/tests/spaces/test_hypercube_graph.py index 8b5d6389..cb5ac9c9 100644 --- a/tests/spaces/test_hypercube_graph.py +++ b/tests/spaces/test_hypercube_graph.py @@ -1,13 +1,11 @@ import lab as B import numpy as np import pytest -from opt_einsum import contract as einsum from plum import Tuple from geometric_kernels.kernels import MaternGeometricKernel from geometric_kernels.spaces import HypercubeGraph -from geometric_kernels.utils.special_functions import hypercube_graph_heat_kernel -from geometric_kernels.utils.utils import binary_vectors_and_subsets, chain +from geometric_kernels.utils.kernel_formulas import hypercube_graph_heat_kernel from ..helper import check_function_with_backend @@ -19,180 +17,39 @@ def inputs(request) -> Tuple[B.Numeric]: - space is a HypercubeGraph object with dimension equal to request.param, - eigenfunctions is the respective Eigenfunctions object with at most 5 levels, - X is a random sample of random size from the space, - - X2 is another random sample of random size from the space. + - X2 is another random sample of random size from the space, + - weights is an array of positive numbers of shape (eigenfunctions.num_levels, 1). """ d = request.param space = HypercubeGraph(d) eigenfunctions = space.get_eigenfunctions(min(space.dim + 1, 5)) - key = np.random.RandomState() + key = np.random.RandomState(0) N, N2 = key.randint(low=1, high=min(2**d, 10) + 1, size=2) key, X = space.random(key, N) key, X2 = space.random(key, N2) - return space, eigenfunctions, X, X2 + # These weights are used for testing the weighted outerproduct, they + # should be positive. + weights = np.random.rand(eigenfunctions.num_levels, 1) ** 2 + 1e-5 - -@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) -def test_call_eigenfunctions(inputs: Tuple[B.NPNumeric, B.NPNumeric], backend): - _, eigenfunctions, X, _ = inputs - - # Check that the eigenfunctions can be called, returning the right type and shape. - check_function_with_backend( - backend, - (X.shape[0], eigenfunctions.num_eigenfunctions), - eigenfunctions, - X, - compare_to_result=lambda res, f_out: f_out.shape == res, - ) + return space, eigenfunctions, X, X2, weights def test_numbers_of_eigenfunctions(inputs): - space, eigenfunctions, _, _ = inputs + space, eigenfunctions, _, _, _ = inputs num_levels = eigenfunctions.num_levels - # Check that the length of the `num_eigenfunctions_per_level` list is correct. - assert len(eigenfunctions.num_eigenfunctions_per_level) == num_levels - # Check that the first eigenspace is 1-dimensional. - assert eigenfunctions.num_eigenfunctions_per_level[0] == 1 + # If the number of levels is maximal, check that the number of # eigenfunctions is equal to the number of binary vectors of size `space.dim`. if num_levels == space.dim + 1: assert eigenfunctions.num_eigenfunctions == 2**space.dim - # Check that dimensions of eigenspaces are always positive. - for i in range(num_levels): - assert eigenfunctions.num_eigenfunctions_per_level[i] > 0 - - num_eigenfunctions_manual = sum(eigenfunctions.num_eigenfunctions_per_level) - # Check that `num_eigenfunctions_per_level` sum up to the total number of - # eigenfunctions. - assert num_eigenfunctions_manual == eigenfunctions.num_eigenfunctions - - -@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) -def test_orthonormality(inputs, backend): - space, _, _, _ = inputs - - if space.dim > 5: - pytest.skip("Test is too slow for dim > 5") - - eigenfunctions = space.get_eigenfunctions(space.dim + 1) - - X, _ = binary_vectors_and_subsets(space.dim) - - # Check that the eigenfunctions are orthonormal with respect to the inner - # product = 2**(-d) sum_i x_i y_i. - check_function_with_backend( - backend, - np.eye(2**space.dim) * 2**space.dim, - lambda X: B.matmul(B.T(eigenfunctions(X)), eigenfunctions(X)), - X, - ) - - -@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) -def test_weighted_outerproduct_with_addition_theorem(inputs, backend): - _, eigenfunctions, X, X2 = inputs - num_levels = eigenfunctions.num_levels - - weights = np.random.rand(num_levels, 1) - chained_weights = chain( - weights.squeeze(), eigenfunctions.num_eigenfunctions_per_level - ) - - Phi_X = eigenfunctions(X) - Phi_X2 = eigenfunctions(X2) - result = einsum("ni,ki,i->nk", Phi_X, Phi_X2, chained_weights) - - # Check that `weighted_outerproduct`, which is based on the addition theorem, - # returns the same result as the direct computation involving individual - # eigenfunctions. - check_function_with_backend( - backend, result, eigenfunctions.weighted_outerproduct, weights, X, X2 - ) - - -@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) -def test_weighted_outerproduct_with_addition_theorem_one_input(inputs, backend): - _, eigenfunctions, X, _ = inputs - num_levels = eigenfunctions.num_levels - - weights = np.random.rand(num_levels, 1) - - result = eigenfunctions.weighted_outerproduct(weights, X, X) - - # Check that `weighted_outerproduct`, when given only X (but not X2), - # uses X2=X. - check_function_with_backend( - backend, - result, - eigenfunctions.weighted_outerproduct, - weights, - X, - ) - - -@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) -def test_weighted_outerproduct_diag(inputs, backend): - _, eigenfunctions, X, _ = inputs - num_levels = eigenfunctions.num_levels - - weights = np.random.rand(num_levels, 1) - - result = np.diag(eigenfunctions.weighted_outerproduct(weights, X, X)) - - # Check that `weighted_outerproduct_diag` returns the same result as the - # diagonal of the full `weighted_outerproduct`. - check_function_with_backend( - backend, - result, - eigenfunctions.weighted_outerproduct_diag, - weights, - X, - ) - - -@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) -def test_weighted_outerproduct_against_phi_product(inputs, backend): - _, eigenfunctions, X, X2 = inputs - num_levels = eigenfunctions.num_levels - - sum_phi_phi_for_level = eigenfunctions.phi_product(X, X2) - - weights = np.random.rand(num_levels, 1) - result = B.einsum("id,...nki->...nk", weights, sum_phi_phi_for_level) - - # Check that `weighted_outerproduct`, which for HypercubeGraph has a - # dedicated implementation, returns the same result as the usual way of - # computing the `weighted_outerproduct` (based on the `phi_product` method). - check_function_with_backend( - backend, result, eigenfunctions.weighted_outerproduct, weights, X, X2 - ) - - -@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) -def test_weighted_outerproduct_diag_against_phi_product(inputs, backend): - _, eigenfunctions, X, _ = inputs - num_levels = eigenfunctions.num_levels - - phi_product_diag = eigenfunctions.phi_product_diag(X) - - weights = np.random.rand(num_levels, 1) - result = B.einsum("id,ni->n", weights, phi_product_diag) # [N,] - - # Check that `weighted_outerproduct_diag`, which for HypercubeGraph has a - # dedicated implementation, returns the same result as the usual way of - # computing the `weighted_outerproduct_diag` (based on the - # `phi_product_diag` method). - check_function_with_backend( - backend, result, eigenfunctions.weighted_outerproduct_diag, weights, X - ) - @pytest.mark.parametrize("lengthscale", [1.0, 5.0, 10.0]) @pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) def test_against_analytic_heat_kernel(inputs, lengthscale, backend): - space, _, X, X2 = inputs + space, _, X, X2, _ = inputs lengthscale = np.array([lengthscale]) result = hypercube_graph_heat_kernel(lengthscale, X, X2) @@ -204,11 +61,8 @@ def test_against_analytic_heat_kernel(inputs, lengthscale, backend): check_function_with_backend( backend, result, - lambda nu, lengthscale, X, X2: kernel.K( - {"nu": nu, "lengthscale": lengthscale}, X, X2 - ), - np.array([np.inf]), - lengthscale, + kernel.K, + {"nu": np.array([np.inf]), "lengthscale": lengthscale}, X, X2, atol=1e-2, diff --git a/tests/spaces/test_hypersphere.py b/tests/spaces/test_hypersphere.py index 860dee00..006469ad 100644 --- a/tests/spaces/test_hypersphere.py +++ b/tests/spaces/test_hypersphere.py @@ -1,115 +1,61 @@ import lab as B import numpy as np -import pytest -from plum import Tuple -from geometric_kernels.spaces.hypersphere import SphericalHarmonics -from geometric_kernels.utils.utils import chain +from geometric_kernels.spaces import Hypersphere -class Consts: - seed = 42 - dimension = 2 - num_data = 7 - num_data2 = 5 - num_eigenfunctions = 9 +def test_sphere_heat_kernel(): + # Tests that the heat kernel on the sphere solves the heat equation. This + # test only uses torch, as lab doesn't support backend-independent autodiff. + import torch + import geometric_kernels.torch # noqa + from geometric_kernels.kernels import MaternKarhunenLoeveKernel + from geometric_kernels.utils.manifold_utils import manifold_laplacian -@pytest.fixture(name="eigenfunctions") -def _eigenfunctions_fixture(): - return SphericalHarmonics(Consts.dimension, Consts.num_eigenfunctions) + _TRUNCATION_LEVEL = 10 + # Parameters + grid_size = 4 + nb_samples = 10 + dimension = 3 -@pytest.fixture(name="inputs") -def _inputs_fixure(request) -> Tuple[B.Numeric]: - def _norm(v): - return np.sum(v**2, axis=-1, keepdims=True) ** 0.5 + # Create manifold + hypersphere = Hypersphere(dim=dimension) - np.random.seed(Consts.seed) - value = np.random.randn(Consts.num_data, Consts.dimension + 1) - value = value / _norm(value) - value2 = np.random.randn(Consts.num_data2, Consts.dimension + 1) - value2 = value2 / _norm(value2) - return value, value2 - - -@pytest.mark.parametrize( - "dim, num, expected_num, expected_num_levels", - [(2, 3, 9, 3), (2, 4, 16, 4), (3, 3, 14, 3), (3, 4, 30, 4), (8, 2, 10, 2)], -) -def test_shape_eigenfunctions(dim, num, expected_num, expected_num_levels): - sph_harmonics = SphericalHarmonics(dim, num) - assert len(sph_harmonics._spherical_harmonics) == sph_harmonics.num_eigenfunctions - assert sph_harmonics.num_eigenfunctions == expected_num - assert sph_harmonics.num_levels == expected_num_levels - - -def test_call_eigenfunctions( - inputs: Tuple[B.Numeric, B.Numeric], - eigenfunctions: SphericalHarmonics, -): - inputs, _ = inputs - output = eigenfunctions(inputs) - assert output.shape == (Consts.num_data, eigenfunctions.num_eigenfunctions) - - -def test_eigenfunctions_shape(eigenfunctions: SphericalHarmonics): - num_eigenfunctions_manual = np.sum(eigenfunctions.num_eigenfunctions_per_level) - assert num_eigenfunctions_manual == eigenfunctions.num_eigenfunctions - assert len(eigenfunctions.num_eigenfunctions_per_level) == eigenfunctions.num_levels - - -def test_weighted_outerproduct_with_addition_theorem( - inputs, eigenfunctions: SphericalHarmonics -): - """ - Eigenfunction will use addition theorem to compute outerproduct. We compare against the - naive implementation. - """ - inputs, inputs2 = inputs - weights_per_level = np.random.randn(eigenfunctions.num_levels) - chained_weights = chain( - weights_per_level, eigenfunctions.num_eigenfunctions_per_level + # Generate samples + ts = torch.linspace(0.1, 1, grid_size, requires_grad=True) + xs = torch.tensor( + np.array(hypersphere.random_point(nb_samples)), requires_grad=True ) - weights = B.reshape(weights_per_level, -1, 1) - actual = B.to_numpy(eigenfunctions.weighted_outerproduct(weights, inputs, inputs2)) - - Phi_X = eigenfunctions(inputs) - Phi_X2 = eigenfunctions(inputs2) - expected = B.einsum("ni,ki,i->nk", Phi_X, Phi_X2, chained_weights) - np.testing.assert_array_almost_equal(actual, expected) - - -def test_weighted_outerproduct_with_addition_theorem_same_input( - inputs, eigenfunctions: SphericalHarmonics -): - """ - Eigenfunction will use addition theorem to compute outerproduct. We compare against the - naive implementation. - """ - inputs, _ = inputs - weights_per_level = np.random.randn(eigenfunctions.num_levels) - weights = B.reshape(weights_per_level, -1, 1) - first = B.to_numpy(eigenfunctions.weighted_outerproduct(weights, inputs, inputs)) - second = B.to_numpy(eigenfunctions.weighted_outerproduct(weights, inputs, None)) - np.testing.assert_array_almost_equal(first, second) - - -def test_weighted_outerproduct_diag_with_addition_theorem( - inputs, eigenfunctions: SphericalHarmonics -): - """ - Eigenfunction will use addition theorem to compute outerproduct. We compare against the - naive implementation. - """ - inputs, _ = inputs - weights_per_level = np.random.randn(eigenfunctions.num_levels) - chained_weights = chain( - weights_per_level, eigenfunctions.num_eigenfunctions_per_level - ) - weights = B.reshape(weights_per_level, -1, 1) - actual = eigenfunctions.weighted_outerproduct_diag(weights, inputs) - - Phi_X = eigenfunctions(inputs) - expected = B.einsum("ni,i->n", Phi_X**2, chained_weights) - np.testing.assert_array_almost_equal(B.to_numpy(actual), B.to_numpy(expected)) + ys = xs + + # Define kernel + kernel = MaternKarhunenLoeveKernel(hypersphere, _TRUNCATION_LEVEL, normalize=False) + params = kernel.init_params() + params["nu"] = torch.tensor([torch.inf]) + + # Define heat kernel function + def heat_kernel(t, x, y): + params["lengthscale"] = B.reshape(B.sqrt(2 * t), 1) + return kernel.K(params, x, y) + + for t in ts: + for x in xs: + for y in ys: + # Compute the derivative of the kernel function wrt t + dfdt, _, _ = torch.autograd.grad( + heat_kernel(t, x[None], y[None]), (t, x, y) + ) + # Compute the Laplacian of the kernel on the manifold + egrad = lambda u: torch.autograd.grad( # noqa + heat_kernel(t, u[None], y[None]), (t, u, y) + )[ + 1 + ] # noqa + fx = lambda u: heat_kernel(t, u[None], y[None]) # noqa + ehess = lambda u, h: torch.autograd.functional.hvp(fx, u, h)[1] # noqa + lapf = manifold_laplacian(x, hypersphere, egrad, ehess) + + # Check that they match + np.testing.assert_allclose(dfdt.detach().numpy(), lapf, atol=1e-3) diff --git a/tests/spaces/test_lie_groups.py b/tests/spaces/test_lie_groups.py index 863580d8..c5a1d327 100644 --- a/tests/spaces/test_lie_groups.py +++ b/tests/spaces/test_lie_groups.py @@ -1,33 +1,36 @@ -import itertools - import lab as B import numpy as np import pytest -from numpy.testing import assert_allclose -from geometric_kernels.feature_maps import RandomPhaseFeatureMapCompact -from geometric_kernels.kernels import MaternKarhunenLoeveKernel +from geometric_kernels.kernels.matern_kernel import default_num +from geometric_kernels.lab_extras import complex_conj from geometric_kernels.spaces import SpecialOrthogonal, SpecialUnitary - -@pytest.fixture(name="group_cls", params=["so", "su"]) -def _group_cls(request): - if request.param == "so": - return SpecialOrthogonal - elif request.param == "su": - return SpecialUnitary +from ..helper import check_function_with_backend, compact_matrix_lie_groups -@pytest.fixture(name="group", params=[3, 5]) -def _group(group_cls, request): - group = group_cls(n=request.param) - return group +@pytest.fixture( + params=compact_matrix_lie_groups(), + ids=str, +) +def inputs(request): + """ + Returns a tuple (space, eigenfunctions, X, X2) where: + - space = request.param, + - eigenfunctions = space.get_eigenfunctions(num_levels), with reasonable num_levels + - X is a random sample of random size from the space, + - X2 is another random sample of random size from the space, + """ + space = request.param + num_levels = min(10, default_num(space)) + eigenfunctions = space.get_eigenfunctions(num_levels) + key = np.random.RandomState(0) + N, N2 = key.randint(low=1, high=100 + 1, size=2) + key, X = space.random(key, N) + key, X2 = space.random(key, N2) -@pytest.fixture(name="group_and_eigf", params=[10]) -def _group_and_eigf(group, request): - eigf = group.get_eigenfunctions(num=request.param) - return group, eigf + return space, eigenfunctions, X, X2 def get_dtype(group): @@ -39,93 +42,90 @@ def get_dtype(group): raise ValueError() -def test_group_inverse(group_and_eigf): - group, eigenfunctions = group_and_eigf - dtype = get_dtype(group) +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_group_inverse(inputs, backend): + group, _, X, _ = inputs - key = B.create_random_state(dtype, seed=0) + result = np.eye(group.n, dtype=get_dtype(group)) + result = np.broadcast_to(result, (X.shape[0], group.n, group.n)) - b1, b2 = 10, 10 - key, x = group.random(key, b1) - key, y = group.random(key, b2) + check_function_with_backend( + backend, + result, + lambda X: B.matmul(X, group.inverse(X)), + X, + ) - eye_ = np.matmul(x, group.inverse(x))[None, ...] - diff = eye_ - np.eye(group.n, dtype=dtype) - zeros = np.zeros_like(eye_) - assert_allclose(diff, zeros, atol=1e-5) +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_character_conj_invariant(inputs, backend): + group, eigenfunctions, X, G = inputs + # Truncate X and G to have the same length + n_xs = min(X.shape[0], G.shape[0]) + X = X[:n_xs, :, :] + G = G[:n_xs, :, :] -def test_character_conj_invariant(group_and_eigf): - group, eigenfunctions = group_and_eigf - dtype = get_dtype(group) + def gammas_diff(X, G, chi): + conjugates = B.matmul(B.matmul(G, X), group.inverse(G)) + conj_gammas = eigenfunctions._torus_representative(conjugates) - key = B.create_random_state(dtype, seed=0) + xs_gammas = eigenfunctions._torus_representative(X) - num_samples_x = 20 - num_samples_g = 20 - key, xs = group.random(key, num_samples_x) - key, gs = group.random(key, num_samples_g) - conjugates = np.matmul(np.matmul(gs, xs), group.inverse(gs)) + return chi(xs_gammas) - chi(conj_gammas) - conj_gammas = eigenfunctions._torus_representative(conjugates) - xs_gammas = eigenfunctions._torus_representative(xs) for chi in eigenfunctions._characters: - chi_vals_xs = chi(xs_gammas) - chi_vals_conj = chi(conj_gammas) - assert_allclose(chi_vals_xs, chi_vals_conj) - - -def test_character_at_identity(group_and_eigf): - group, eigenfunctions = group_and_eigf - dtype = get_dtype(group) - - identity = np.eye(group.n, dtype=dtype).reshape(1, group.n, group.n) - identity_gammas = eigenfunctions._torus_representative(identity) - dimensions = eigenfunctions._dimensions - characters = eigenfunctions._characters - for chi, dim in zip(characters, dimensions): - chi_val = chi(identity_gammas) - assert_allclose(chi_val.real, dim) - assert_allclose(chi_val.imag, 0) - - -def test_characters_orthogonal(group_and_eigf): - group, eigenfunctions = group_and_eigf - dtype = get_dtype(group) - order = eigenfunctions.num_levels - - key = B.create_random_state(dtype, seed=0) - - num_samples_x = 5 * 10**5 - key, xs = group.random(key, num_samples_x) - gammas = eigenfunctions._torus_representative(xs) - characters = eigenfunctions._characters - scalar_products = np.zeros((order, order), dtype=dtype) - for a, b in itertools.product(enumerate(characters), repeat=2): - i, chi1 = a - j, chi2 = b - scalar_products[i, j] = np.mean((np.conj(chi1(gammas)) * chi2(gammas)).real) - - assert_allclose(scalar_products, np.eye(order, dtype=dtype), atol=5e-2) - - -def test_feature_map(group_and_eigf): - group, eigenfunctions = group_and_eigf - order = eigenfunctions.num_levels - dtype = get_dtype(group) - key = B.create_random_state(dtype, seed=0) - - kernel = MaternKarhunenLoeveKernel(group, order, normalize=True) - param = dict(lengthscale=np.array([10]), nu=np.array([1.5])) - - feature_order = 5000 - feature_map = RandomPhaseFeatureMapCompact(group, order, feature_order) - - key, x = group.random(key, 10) - - K_xx = (kernel.K(param, x, x)).real - key, embed_x = feature_map(x, param, key=key, normalize=True) - F_xx = (B.einsum("ni,mi-> nm", embed_x, embed_x.conj())).real - - assert_allclose(K_xx, F_xx, atol=5e-2) + check_function_with_backend( + backend, + np.zeros((n_xs,)), + lambda X, G: gammas_diff(X, G, chi), + X, + G, + atol=1e-3, + ) + + +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_character_at_identity(inputs, backend): + group, eigenfunctions, _, _ = inputs + + for chi, dim in zip(eigenfunctions._characters, eigenfunctions._dimensions): + check_function_with_backend( + backend, + np.array([dim], dtype=get_dtype(group)), + lambda X: B.real(chi(eigenfunctions._torus_representative(X))), + np.eye(group.n, dtype=get_dtype(group))[None, ...], + ) + + check_function_with_backend( + backend, + np.array([0], dtype=get_dtype(group)), + lambda X: B.imag(chi(eigenfunctions._torus_representative(X))), + np.eye(group.n, dtype=get_dtype(group))[None, ...], + ) + + +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_characters_orthogonal(inputs, backend): + group, eigenfunctions, _, _ = inputs + + num_samples = 10000 + key = np.random.RandomState(0) + _, X = group.random(key, num_samples) + + def all_char_vals(X): + gammas = eigenfunctions._torus_representative(X) + values = [ + chi(gammas)[..., None] # [num_samples, 1] + for chi in eigenfunctions._characters + ] + + return B.concat(*values, axis=-1) + + check_function_with_backend( + backend, + np.eye(eigenfunctions.num_levels, dtype=get_dtype(group)), + lambda X: complex_conj(B.T(all_char_vals(X))) @ all_char_vals(X) / num_samples, + X, + atol=0.4, # very loose, but helps make sure the diagonal is close to 1 while the rest is close to 0 + ) diff --git a/tests/spaces/test_mesh.py b/tests/spaces/test_mesh.py index 5f26dd71..164bcc6e 100644 --- a/tests/spaces/test_mesh.py +++ b/tests/spaces/test_mesh.py @@ -1,14 +1,13 @@ from pathlib import Path import numpy as np -from numpy.testing import assert_array_almost_equal -from pytest import fixture +import pytest from geometric_kernels.spaces import Mesh -@fixture(name="mesh") -def fixture_get_mesh() -> Mesh: +@pytest.fixture() +def mesh() -> Mesh: filename = Path(__file__).parent / "../teddy.obj" mesh = Mesh.load_mesh(str(filename)) return mesh @@ -17,7 +16,7 @@ def fixture_get_mesh() -> Mesh: def test_mesh_shapes(): Nv = 11 # num vertices Nf = 13 # num faces - dim = 3 # dimension + dim = 3 # ambient dimension vertices = np.random.randn(Nv, dim) faces = np.random.randint(0, Nv, size=(Nf, 3)) mesh = Mesh(vertices=vertices, faces=faces) @@ -39,8 +38,3 @@ def test_eigenvectors(mesh: Mesh): assert mesh.get_eigenvectors(10).shape == (mesh.num_vertices, 10) assert mesh.get_eigenvectors(13).shape == (mesh.num_vertices, 13) assert set(mesh.cache.keys()) == set([10, 13]) - - -def test_orthonormality_eigenvectors(mesh: Mesh): - evecs = mesh.get_eigenvectors(10) # [Nv, 10] - assert_array_almost_equal(evecs.T @ evecs, mesh.num_vertices * np.eye(10)) diff --git a/tests/spaces/test_product_discrete_spectrum_space.py b/tests/spaces/test_product_discrete_spectrum_space.py new file mode 100644 index 00000000..5fada669 --- /dev/null +++ b/tests/spaces/test_product_discrete_spectrum_space.py @@ -0,0 +1,56 @@ +import numpy as np +import pytest + +from geometric_kernels.kernels import MaternGeometricKernel +from geometric_kernels.spaces import ( + Circle, + ProductDiscreteSpectrumSpace, + SpecialUnitary, +) +from geometric_kernels.utils.product import make_product + +from ..helper import check_function_with_backend + +_NUM_LEVELS = 20 + + +@pytest.mark.parametrize( + "factor1, factor2", [(Circle(), Circle()), (Circle(), SpecialUnitary(2))], ids=str +) +@pytest.mark.parametrize("lengthscale", [0.1, 0.5, 1.0, 2.0, 5.0]) +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_heat_kernel_is_product_of_heat_kernels(factor1, factor2, lengthscale, backend): + product = ProductDiscreteSpectrumSpace( + factor1, factor2, num_levels=_NUM_LEVELS**2, num_levels_per_space=_NUM_LEVELS + ) + + key = np.random.RandomState(0) + key, xs_factor1 = factor1.random(key, 10) + key, xs_factor2 = factor2.random(key, 10) + + kernel_product = MaternGeometricKernel(product, num=_NUM_LEVELS**2) + kernel_factor1 = MaternGeometricKernel(factor1, num=_NUM_LEVELS) + kernel_factor2 = MaternGeometricKernel(factor2, num=_NUM_LEVELS) + + def K_diff(nu, lengthscale, xs_factor1, xs_factor2): + params = {"nu": nu, "lengthscale": lengthscale} + + xs_product = make_product([xs_factor1, xs_factor2]) + + K_product = kernel_product.K(params, xs_product, xs_product) + K_factor1 = kernel_factor1.K(params, xs_factor1, xs_factor1) + K_factor2 = kernel_factor2.K(params, xs_factor2, xs_factor2) + + return K_product - K_factor1 * K_factor2 + + # Check that the heat kernel on the ProductDiscreteSpectrumSpace coincides + # with the product of heat kernels on the factors. + check_function_with_backend( + backend, + np.zeros((xs_factor1.shape[0], xs_factor2.shape[0])), + K_diff, + np.array([np.inf]), + np.array([lengthscale]), + xs_factor1, + xs_factor2, + ) diff --git a/tests/spaces/test_spd.py b/tests/spaces/test_spd.py new file mode 100644 index 00000000..9e4729e2 --- /dev/null +++ b/tests/spaces/test_spd.py @@ -0,0 +1,45 @@ +import lab as B +import numpy as np +import pytest + +from geometric_kernels.kernels import MaternGeometricKernel +from geometric_kernels.spaces import SymmetricPositiveDefiniteMatrices +from geometric_kernels.utils.kernel_formulas import spd_heat_kernel_2x2 + +from ..helper import check_function_with_backend, create_random_state + + +@pytest.mark.parametrize("lengthscale", [2.0]) +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_equivalence_kernel(lengthscale, backend): + space = SymmetricPositiveDefiniteMatrices(2) + + key = np.random.RandomState(0) + key, X = space.random(key, 5) + X2 = X.copy() + + t = lengthscale * lengthscale / 2 + result = spd_heat_kernel_2x2(t, X, X2) + + kernel = MaternGeometricKernel(space, key=create_random_state(backend)) + + def compare_to_result(res, f_out): + return ( + np.linalg.norm(res - B.to_numpy(f_out)) + / np.sqrt(res.shape[0] * res.shape[1]) + < 1e-1 + ) + + # Check that MaternGeometricKernel on SymmetricPositiveDefiniteMatrices(2) + # with nu=inf coincides with the semi-analytic formula from :cite:t:`sawyer1992`. + # We are checking the equivalence on average, computing the norm between + # the two covariance matrices. + check_function_with_backend( + backend, + result, + kernel.K, + {"nu": np.array([np.inf]), "lengthscale": np.array([lengthscale])}, + X, + X2, + compare_to_result=compare_to_result, + ) diff --git a/tests/test_dtypes.py b/tests/test_dtypes.py deleted file mode 100644 index 5d347744..00000000 --- a/tests/test_dtypes.py +++ /dev/null @@ -1,213 +0,0 @@ -import sys - -import jax.numpy as jnp -import lab as B -import numpy as np -import pytest -import tensorflow as tf -import torch - -from geometric_kernels.feature_maps import ( - DeterministicFeatureMapCompact, - RandomPhaseFeatureMapCompact, - RandomPhaseFeatureMapNoncompact, - RejectionSamplingFeatureMapHyperbolic, - RejectionSamplingFeatureMapSPD, -) -from geometric_kernels.kernels.karhunen_loeve import MaternKarhunenLoeveKernel -from geometric_kernels.spaces import ( - Circle, - Hyperbolic, - Hypersphere, - Mesh, - SymmetricPositiveDefiniteMatrices, -) - - -def to_typed_ndarray(value, dtype): - if dtype == "float32": - return value.astype(np.float32) - elif dtype == "float64": - return value.astype(np.float64) - else: - raise ValueError("Unknown dtype: {}".format(dtype)) - - -def to_typed_tensor(value, backend): - if backend == "tensorflow": - return tf.convert_to_tensor(value) - elif backend in ["torch", "pytorch"]: - return torch.tensor(value) - elif backend == "numpy": - return value - elif backend == "jax": - return jnp.array(value) - else: - raise ValueError("Unknown backend: {}".format(backend)) - - -def mesh_point(): - n_base = 4 - n_vertices = 2 * n_base - vertices = np.array( - [ - ( - 1.0 * (i % 2), - np.cos(2 * np.pi * (i // 2) / n_base), - np.sin(2 * np.pi * (i // 2) / n_base), - ) - for i in range(n_vertices) - ] - ) - faces = np.array( - [ - (i % n_vertices, (i + 1) % n_vertices, (i + 2) % n_vertices) - for i in range(n_vertices) - ] # box without sides - + [ - (i % 2, (i + 2) % n_vertices, (i + 4) % n_vertices) - for i in range(n_vertices - 4) - ] # sides - ) - # this is just a box - - mesh = Mesh(vertices, faces) - point = np.array([0]).reshape(1, 1) - - return mesh, point - - -def circle_point(): - circle = Circle() - - point = np.array([0]).reshape(1, 1) - - return circle, point - - -def hypersphere_point(): - hypersphere = Hypersphere(dim=2) - - point = hypersphere.random_point(1).reshape(1, -1) - - return hypersphere, point - - -def hyperbolic_point(): - hyperboloid = Hyperbolic(dim=2) - - point = hyperboloid.random_point(1).reshape(1, -1) - - return hyperboloid, point - - -def spd_point(): - spd = SymmetricPositiveDefiniteMatrices(2) - - point = spd.random_point(1).reshape(1, 2, 2) - - return spd, point - - -@pytest.fixture(name="noncompact_spacepoint", params=["hyperbolic", "spd"]) -def _noncompact_spacepoint(request): - if request.param == "hyperbolic": - return hyperbolic_point() - elif request.param == "spd": - return spd_point() - else: - raise ValueError("Unknown space {}".format(request.param)) - - -@pytest.fixture(name="kl_spacepoint", params=["circle", "hypersphere", "mesh"]) -def _kl_spacepoint_fixture(request): - if request.param == "circle": - return circle_point() - elif request.param == "hypersphere": - return hypersphere_point() - elif request.param == "mesh": - return mesh_point() - else: - raise ValueError("Unknown space {}".format(request.param)) - - -@pytest.mark.parametrize("dtype", ["float64", "float32"]) -@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) -def test_karhunen_loeve_dtype(kl_spacepoint, dtype, backend): - space, point = kl_spacepoint - point = to_typed_ndarray(point, dtype) - point = to_typed_tensor(point, backend) - - kernel = MaternKarhunenLoeveKernel(space, 3) - - params = kernel.init_params() - params["nu"] = to_typed_tensor(to_typed_ndarray(np.r_[0.5], dtype), backend) - params["lengthscale"] = to_typed_tensor( - to_typed_ndarray(np.r_[0.5], dtype), backend - ) - - # make sure that it just runs - kernel.K(params, point) - - -@pytest.mark.parametrize("dtype", ["float32", "float64"]) -@pytest.mark.parametrize("backend", ["numpy", "jax", "torch", "tensorflow"]) -def test_feature_map_dtype(kl_spacepoint, dtype, backend): - space, point = kl_spacepoint - point = to_typed_ndarray(point, dtype) - point = to_typed_tensor(point, backend) - - num_levels = 3 - kernel = MaternKarhunenLoeveKernel(space, num_levels) - - params = kernel.init_params() - params["nu"] = to_typed_tensor(to_typed_ndarray(np.r_[0.5], dtype), backend) - params["lengthscale"] = to_typed_tensor( - to_typed_ndarray(np.r_[0.5], dtype), backend - ) - - # make sure it runs - feature_map = DeterministicFeatureMapCompact(space, num_levels) - feature_map(point, params) - - # make sure it runs - key = B.create_random_state(B.dtype(point), seed=1234) - feature_map = RandomPhaseFeatureMapCompact(space, num_levels) - feature_map(point, params, key=key) - - -@pytest.fixture(params=["naive", "rs"]) -def feature_map_on_noncompact(request, noncompact_spacepoint): - space = noncompact_spacepoint[0] - if request.param == "naive": - feature_map = RandomPhaseFeatureMapNoncompact(space, 10) - elif request.param == "rs" and isinstance(space, Hyperbolic): - feature_map = RejectionSamplingFeatureMapHyperbolic(space, 10) - elif request.param == "rs" and isinstance(space, SymmetricPositiveDefiniteMatrices): - feature_map = RejectionSamplingFeatureMapSPD(space, 10) - else: - raise ValueError(f"Unknown feature map {request.param}") - return noncompact_spacepoint + (feature_map,) - - -@pytest.mark.skipif( - sys.version_info < (3, 8), - reason="requires newer numpy version, unavailable in Python<=3.7", -) -@pytest.mark.parametrize("dtype", ["float32", "float64"]) -@pytest.mark.parametrize("backend", ["numpy", "jax", "torch", "tensorflow"]) -@pytest.mark.parametrize("nu", [0.5, np.inf]) -def test_feature_map_noncompact_dtype(feature_map_on_noncompact, dtype, backend, nu): - space, point, feature_map = feature_map_on_noncompact - point = to_typed_ndarray(point, dtype) - point = to_typed_tensor(point, backend) - - params = {} - params["nu"] = to_typed_tensor(to_typed_ndarray(np.r_[nu], dtype), backend) - params["lengthscale"] = to_typed_tensor( - to_typed_ndarray(np.r_[0.5], dtype), backend - ) - - # make sure it runs - key = B.create_random_state(B.dtype(point), seed=1234) - feature_map(point, params, key=key) diff --git a/tests/test_import.py b/tests/test_import.py deleted file mode 100644 index 0b009d43..00000000 --- a/tests/test_import.py +++ /dev/null @@ -1,2 +0,0 @@ -def test_import(): - assert True diff --git a/tests/utils/test_kernel_formulas.py b/tests/utils/test_kernel_formulas.py new file mode 100644 index 00000000..427acfc9 --- /dev/null +++ b/tests/utils/test_kernel_formulas.py @@ -0,0 +1,68 @@ +from math import log, tanh + +import numpy as np +import pytest +from sklearn.metrics.pairwise import rbf_kernel + +from geometric_kernels.spaces import HypercubeGraph +from geometric_kernels.utils.kernel_formulas import hypercube_graph_heat_kernel + +from ..helper import check_function_with_backend + + +@pytest.mark.parametrize("d", [1, 5, 10]) +@pytest.mark.parametrize("lengthscale", [1.0, 5.0, 10.0]) +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_hypercube_graph_heat_kernel(d, lengthscale, backend): + space = HypercubeGraph(d) + + key = np.random.RandomState(0) + N, N2 = key.randint(low=1, high=min(2**d, 10) + 1, size=2) + key, X = space.random(key, N) + key, X2 = space.random(key, N2) + + gamma = -log(tanh(lengthscale**2 / 2)) + result = rbf_kernel(X, X2, gamma=gamma) + + def heat_kernel(lengthscale, X, X2): + return hypercube_graph_heat_kernel( + lengthscale, X, X2, normalized_laplacian=False + ) + + # Checks that the heat kernel on the hypercube graph coincides with the RBF + # restricted onto binary vectors, with appropriately redefined length scale. + check_function_with_backend( + backend, + result, + heat_kernel, + np.array([lengthscale]), + X, + X2, + atol=1e-2, + ) + + if d > 5: + X_first = X[0:1, :3] + X2_first = X2[0:1, :3] + X_second = X[0:1, 3:] + X2_second = X2[0:1, 3:] + + K_first = hypercube_graph_heat_kernel( + np.array([lengthscale]), X_first, X2_first, normalized_laplacian=False + ) + K_second = hypercube_graph_heat_kernel( + np.array([lengthscale]), X_second, X2_second, normalized_laplacian=False + ) + + result = K_first * K_second + + # Checks that the heat kernel of the product is equal to the product + # of heat kernels. + check_function_with_backend( + backend, + result, + heat_kernel, + np.array([lengthscale]), + X[0:1, :], + X2[0:1, :], + ) diff --git a/tests/utils/test_manifold_utils.py b/tests/utils/test_manifold_utils.py new file mode 100644 index 00000000..08020cd9 --- /dev/null +++ b/tests/utils/test_manifold_utils.py @@ -0,0 +1,32 @@ +import numpy as np +import pytest + +from geometric_kernels.spaces import Hyperbolic +from geometric_kernels.utils.manifold_utils import hyperbolic_distance + +from ..helper import check_function_with_backend + + +@pytest.mark.parametrize("dim", [2, 3, 9, 10]) +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_hyperboloid_distance(dim, backend): + space = Hyperbolic(dim=dim) + + key = np.random.RandomState(0) + N, N2 = key.randint(low=2, high=15, size=2) + key, X = space.random(key, N) + key, X2 = space.random(key, N2) + + X_expanded = np.tile(X[..., None, :], (1, X2.shape[0], 1)) # (N, M, n+1) + X2_expanded = np.tile(X2[None], (X.shape[0], 1, 1)) # (N, M, n+1) + result = space.metric.dist(X_expanded, X2_expanded) + + # Check that our implementation of the hyperbolic distance coincides with + # the one from geomstats. + check_function_with_backend( + backend, + result, + hyperbolic_distance, + X, + X2, + ) diff --git a/tests/utils/test_special_functions.py b/tests/utils/test_special_functions.py index 02ccba7c..72f8f398 100644 --- a/tests/utils/test_special_functions.py +++ b/tests/utils/test_special_functions.py @@ -1,13 +1,10 @@ -from math import comb, log, tanh +from math import comb import lab as B import numpy as np import pytest -from sklearn.metrics.pairwise import rbf_kernel -from geometric_kernels.spaces import HypercubeGraph from geometric_kernels.utils.special_functions import ( - hypercube_graph_heat_kernel, kravchuk_normalized, walsh_function, ) @@ -69,13 +66,15 @@ def test_kravchuk_polynomials(all_xs_and_combs, backend): keepdims=True, ) + def krav(x0, X): + return comb(d, j) * kravchuk_normalized(d, j, hamming_distance(x0, X)) + # Checks that Kravchuk polynomials coincide with certain sums of # the Walsh functions. check_function_with_backend( backend, result, - lambda x0, X: comb(d, j) - * kravchuk_normalized(d, j, hamming_distance(x0, X)), + krav, x0, X, ) @@ -94,14 +93,15 @@ def test_kravchuk_precomputed(all_xs_and_combs, backend): cur_kravchuk_normalized = kravchuk_normalized(d, j, hamming_distance(x0, X)) + def krav(x0, X, kn1, kn2): + return kravchuk_normalized(d, j, hamming_distance(x0, X), kn1, kn2) + # Checks that Kravchuk polynomials coincide with certain sums of # the Walsh functions. check_function_with_backend( backend, cur_kravchuk_normalized, - lambda x0, X, kn1, kn2: kravchuk_normalized( - d, j, hamming_distance(x0, X), kn1, kn2 - ), + krav, x0, X, kravchuk_normalized_j_minus_1, @@ -110,60 +110,3 @@ def test_kravchuk_precomputed(all_xs_and_combs, backend): kravchuk_normalized_j_minus_2 = kravchuk_normalized_j_minus_1 kravchuk_normalized_j_minus_1 = cur_kravchuk_normalized - - -@pytest.mark.parametrize("d", [1, 5, 10]) -@pytest.mark.parametrize("lengthscale", [1.0, 5.0, 10.0]) -@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) -def test_hypercube_graph_heat_kernel(d, lengthscale, backend): - space = HypercubeGraph(d) - - key = np.random.RandomState() - N, N2 = key.randint(low=1, high=min(2**d, 10) + 1, size=2) - key, X = space.random(key, N) - key, X2 = space.random(key, N2) - - gamma = -log(tanh(lengthscale**2 / 2)) - result = rbf_kernel(X, X2, gamma=gamma) - - # Checks that the heat kernel on the hypercube graph coincides with the RBF - # restricted onto binary vectors, with appropriately redefined length scale. - check_function_with_backend( - backend, - result, - lambda lengthscale, X, X2: hypercube_graph_heat_kernel( - lengthscale, X, X2, normalized_laplacian=False - ), - np.array([lengthscale]), - X, - X2, - atol=1e-2, - ) - - if d > 5: - X_first = X[0:1, :3] - X2_first = X2[0:1, :3] - X_second = X[0:1, 3:] - X2_second = X2[0:1, 3:] - - K_first = hypercube_graph_heat_kernel( - np.array([lengthscale]), X_first, X2_first, normalized_laplacian=False - ) - K_second = hypercube_graph_heat_kernel( - np.array([lengthscale]), X_second, X2_second, normalized_laplacian=False - ) - - result = K_first * K_second - - # Checks that the heat kernel of the product is equal to the product - # of heat kernels. - check_function_with_backend( - backend, - result, - lambda lengthscale, X, X2: hypercube_graph_heat_kernel( - lengthscale, X, X2, normalized_laplacian=False - ), - np.array([lengthscale]), - X[0:1, :], - X2[0:1, :], - )