From 66a591f65e7dfc8606cac70443a362275fc620d8 Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Wed, 11 Sep 2024 09:30:31 +0100 Subject: [PATCH 01/35] In test_hypercube_graph, add random weights to a fixture and make them positive. --- tests/spaces/test_hypercube_graph.py | 39 ++++++++++++---------------- 1 file changed, 16 insertions(+), 23 deletions(-) diff --git a/tests/spaces/test_hypercube_graph.py b/tests/spaces/test_hypercube_graph.py index 8b5d6389..761fe614 100644 --- a/tests/spaces/test_hypercube_graph.py +++ b/tests/spaces/test_hypercube_graph.py @@ -19,7 +19,8 @@ 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) @@ -30,12 +31,16 @@ def inputs(request) -> Tuple[B.Numeric]: 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 + 0.01 + + return space, eigenfunctions, X, X2, weights @pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) def test_call_eigenfunctions(inputs: Tuple[B.NPNumeric, B.NPNumeric], backend): - _, eigenfunctions, X, _ = inputs + _, eigenfunctions, X, _, _ = inputs # Check that the eigenfunctions can be called, returning the right type and shape. check_function_with_backend( @@ -48,7 +53,7 @@ def test_call_eigenfunctions(inputs: Tuple[B.NPNumeric, B.NPNumeric], backend): 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 @@ -71,7 +76,7 @@ def test_numbers_of_eigenfunctions(inputs): @pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) def test_orthonormality(inputs, backend): - space, _, _, _ = inputs + space, _, _, _, _ = inputs if space.dim > 5: pytest.skip("Test is too slow for dim > 5") @@ -92,10 +97,8 @@ def test_orthonormality(inputs, backend): @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 + _, eigenfunctions, X, X2, weights = inputs - weights = np.random.rand(num_levels, 1) chained_weights = chain( weights.squeeze(), eigenfunctions.num_eigenfunctions_per_level ) @@ -114,10 +117,7 @@ def test_weighted_outerproduct_with_addition_theorem(inputs, backend): @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) + _, eigenfunctions, X, _, weights = inputs result = eigenfunctions.weighted_outerproduct(weights, X, X) @@ -134,10 +134,7 @@ def test_weighted_outerproduct_with_addition_theorem_one_input(inputs, backend): @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) + _, eigenfunctions, X, _, weights = inputs result = np.diag(eigenfunctions.weighted_outerproduct(weights, X, X)) @@ -154,12 +151,10 @@ def test_weighted_outerproduct_diag(inputs, backend): @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 + _, eigenfunctions, X, X2, weights = inputs 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 @@ -172,12 +167,10 @@ def test_weighted_outerproduct_against_phi_product(inputs, backend): @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 + _, eigenfunctions, X, _, weights = inputs 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 @@ -192,7 +185,7 @@ def test_weighted_outerproduct_diag_against_phi_product(inputs, backend): @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) From e1d96dc7ecb69d20f4c2bcaf81304a3f947df551 Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Wed, 11 Sep 2024 11:44:06 +0100 Subject: [PATCH 02/35] Move tests around to align with the current package structure better --- tests/{sampling => feature_maps}/__init__.py | 0 .../test_feature_maps.py | 0 .../test_student_t_sample.py | 0 tests/kernels/test_sphere_heat_kernel.py | 58 ------------------- tests/spaces/test_hypersphere.py | 56 ++++++++++++++++++ tests/test_import.py | 2 - tests/{kernels => }/test_product.py | 0 7 files changed, 56 insertions(+), 60 deletions(-) rename tests/{sampling => feature_maps}/__init__.py (100%) rename tests/{kernels => feature_maps}/test_feature_maps.py (100%) rename tests/{sampling => feature_maps}/test_student_t_sample.py (100%) delete mode 100644 tests/kernels/test_sphere_heat_kernel.py delete mode 100644 tests/test_import.py rename tests/{kernels => }/test_product.py (100%) diff --git a/tests/sampling/__init__.py b/tests/feature_maps/__init__.py similarity index 100% rename from tests/sampling/__init__.py rename to tests/feature_maps/__init__.py diff --git a/tests/kernels/test_feature_maps.py b/tests/feature_maps/test_feature_maps.py similarity index 100% rename from tests/kernels/test_feature_maps.py rename to tests/feature_maps/test_feature_maps.py diff --git a/tests/sampling/test_student_t_sample.py b/tests/feature_maps/test_student_t_sample.py similarity index 100% rename from tests/sampling/test_student_t_sample.py rename to tests/feature_maps/test_student_t_sample.py 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/spaces/test_hypersphere.py b/tests/spaces/test_hypersphere.py index 860dee00..61f51bfd 100644 --- a/tests/spaces/test_hypersphere.py +++ b/tests/spaces/test_hypersphere.py @@ -3,6 +3,7 @@ import pytest from plum import Tuple +from geometric_kernels.spaces import Hypersphere from geometric_kernels.spaces.hypersphere import SphericalHarmonics from geometric_kernels.utils.utils import chain @@ -113,3 +114,58 @@ def test_weighted_outerproduct_diag_with_addition_theorem( 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 test_sphere_heat_kernel(): + import torch + + import geometric_kernels.torch # noqa + from geometric_kernels.kernels import MaternKarhunenLoeveKernel + from geometric_kernels.utils.manifold_utils import manifold_laplacian + + _TRUNCATION_LEVEL = 10 + + # 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/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/kernels/test_product.py b/tests/test_product.py similarity index 100% rename from tests/kernels/test_product.py rename to tests/test_product.py From cc85678331f720e0f3e3120df8d638bf6dc16ce0 Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Fri, 18 Oct 2024 15:44:43 +0200 Subject: [PATCH 03/35] Some fixes and improvements inspired by revising tests. Detailed list below. - Improve lab_extras.take_along_axis - SpecialOrthogonal/SpecialUnitary fixes for non-numpy backends - Add str for all spaces - Hypersphere.num_eigenfunctions refactor to avoid code duplication - Add num_eigenfunctions caching in WeylAdditionTheorem, similar to how it is done in SphericalHarmonics --- .../lab_extras/tensorflow/extras.py | 2 +- geometric_kernels/lab_extras/torch/extras.py | 2 +- geometric_kernels/spaces/circle.py | 3 ++ geometric_kernels/spaces/eigenfunctions.py | 2 + geometric_kernels/spaces/graph.py | 3 ++ geometric_kernels/spaces/hyperbolic.py | 3 ++ geometric_kernels/spaces/hypercube_graph.py | 3 ++ geometric_kernels/spaces/hypersphere.py | 8 ++-- geometric_kernels/spaces/lie_groups.py | 9 ++-- geometric_kernels/spaces/mesh.py | 3 ++ geometric_kernels/spaces/product.py | 3 ++ geometric_kernels/spaces/so.py | 44 +++++++++++++++---- geometric_kernels/spaces/spd.py | 3 ++ geometric_kernels/spaces/su.py | 3 ++ 14 files changed, 73 insertions(+), 18 deletions(-) diff --git a/geometric_kernels/lab_extras/tensorflow/extras.py b/geometric_kernels/lab_extras/tensorflow/extras.py index 8e613f15..ec978c05 100644 --- a/geometric_kernels/lab_extras/tensorflow/extras.py +++ b/geometric_kernels/lab_extras/tensorflow/extras.py @@ -13,7 +13,7 @@ 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) + return tf.experimental.numpy.take_along_axis(a, index, axis=axis) @dispatch diff --git a/geometric_kernels/lab_extras/torch/extras.py b/geometric_kernels/lab_extras/torch/extras.py index 15860ec5..823e563a 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 diff --git a/geometric_kernels/spaces/circle.py b/geometric_kernels/spaces/circle.py index 4edee885..b88f448e 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: """ 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..bbf92f4d 100644 --- a/geometric_kernels/spaces/graph.py +++ b/geometric_kernels/spaces/graph.py @@ -61,6 +61,9 @@ def __init__(self, adjacency_matrix: B.Numeric, normalize_laplacian: bool = Fals self._checks(adjacency_matrix) self._set_laplacian(adjacency_matrix, normalize_laplacian) # type: ignore + def __str__(self): + return f"Graph({self.num_vertices})" + @staticmethod def _checks(adjacency): """ diff --git a/geometric_kernels/spaces/hyperbolic.py b/geometric_kernels/spaces/hyperbolic.py index 61b57ebc..4e0a336f 100644 --- a/geometric_kernels/spaces/hyperbolic.py +++ b/geometric_kernels/spaces/hyperbolic.py @@ -53,6 +53,9 @@ class Hyperbolic(NoncompactSymmetricSpace, gs.geometry.hyperboloid.Hyperboloid): def __init__(self, dim=2): super().__init__(dim=dim) + def __str__(self): + return f"Hyperbolic({self.dimension})" + @property def dimension(self) -> int: """ diff --git a/geometric_kernels/spaces/hypercube_graph.py b/geometric_kernels/spaces/hypercube_graph.py index 0f58b7a6..1c315523 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: """ diff --git a/geometric_kernels/spaces/hypersphere.py b/geometric_kernels/spaces/hypersphere.py index f1674102..39c22b85 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: """ diff --git a/geometric_kernels/spaces/lie_groups.py b/geometric_kernels/spaces/lie_groups.py index 9a79efa4..6e3877e8 100644 --- a/geometric_kernels/spaces/lie_groups.py +++ b/geometric_kernels/spaces/lie_groups.py @@ -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, ...]]: @@ -202,8 +203,8 @@ def _difference(self, X: B.Numeric, X2: B.Numeric) -> B.Numeric: 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 @@ -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..a1cb37e5 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 diff --git a/geometric_kernels/spaces/product.py b/geometric_kernels/spaces/product.py index f278be7a..0d99920b 100644 --- a/geometric_kernels/spaces/product.py +++ b/geometric_kernels/spaces/product.py @@ -440,6 +440,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: """ diff --git a/geometric_kernels/spaces/so.py b/geometric_kernels/spaces/so.py index 779d275a..20397b0b 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, @@ -125,7 +133,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) @@ -143,17 +151,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.real(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: @@ -203,7 +224,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 @@ -245,6 +268,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: """ diff --git a/geometric_kernels/spaces/spd.py b/geometric_kernels/spaces/spd.py index 0c329bd2..7cd9bb34 100644 --- a/geometric_kernels/spaces/spd.py +++ b/geometric_kernels/spaces/spd.py @@ -54,6 +54,9 @@ class SymmetricPositiveDefiniteMatrices( def __init__(self, n): super().__init__(n) + def __str__(self): + return f"SpecialOrthogonal({self.n})" + @property def dimension(self) -> int: """ diff --git a/geometric_kernels/spaces/su.py b/geometric_kernels/spaces/su.py index a6c89d23..5c8dcada 100644 --- a/geometric_kernels/spaces/su.py +++ b/geometric_kernels/spaces/su.py @@ -175,6 +175,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: """ From 623fb6307278fdfe1136c5b206d01f188f5bce6c Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Fri, 18 Oct 2024 16:26:26 +0200 Subject: [PATCH 04/35] Fix take_along_axis for Python 3.9 and below --- geometric_kernels/lab_extras/tensorflow/extras.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/geometric_kernels/lab_extras/tensorflow/extras.py b/geometric_kernels/lab_extras/tensorflow/extras.py index ec978c05..066732e2 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.experimental.numpy.take_along_axis(a, 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 From e3ff505245906981c3c304ac9738e31b53c6c56e Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Mon, 21 Oct 2024 16:21:33 +0200 Subject: [PATCH 05/35] Unify most tests for the Circle, HypercubeGraph, Hypersphere, SpecialOrthogonal, SpecialUnitary spaces buy introducing the new tests/spaces/test_eigenfunctions_basics.py and more: MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - fix a bug in SpecialUnitary under TensorFlow - add Euclidean Matérn -1/2 -3/2 -5/2 and RBF kernels to geometric_kernels/utils/special_functions.py --- geometric_kernels/spaces/su.py | 8 +- geometric_kernels/utils/special_functions.py | 93 ++++++++ tests/spaces/test_circle.py | 199 ++++------------- tests/spaces/test_eigenfunctions_basics.py | 220 +++++++++++++++++++ tests/spaces/test_hypercube_graph.py | 119 +--------- tests/spaces/test_hypersphere.py | 135 +++--------- 6 files changed, 396 insertions(+), 378 deletions(-) create mode 100644 tests/spaces/test_eigenfunctions_basics.py diff --git a/geometric_kernels/spaces/su.py b/geometric_kernels/spaces/su.py index 35479c13..f529765c 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, @@ -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 diff --git a/geometric_kernels/utils/special_functions.py b/geometric_kernels/utils/special_functions.py index 1e27f5b5..69add35c 100644 --- a/geometric_kernels/utils/special_functions.py +++ b/geometric_kernels/utils/special_functions.py @@ -152,3 +152,96 @@ def hypercube_graph_heat_kernel( hamming_distances = B.cast(float_like(X), hamming_distance(X, X2)) return B.tanh(lengthscale**2 / 2) ** hamming_distances + + +def 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 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 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 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/tests/spaces/test_circle.py b/tests/spaces/test_circle.py index 426ff2a5..665bdebf 100644 --- a/tests/spaces/test_circle.py +++ b/tests/spaces/test_circle.py @@ -1,175 +1,72 @@ 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.special_functions import ( + matern_12_kernel, + matern_32_kernel, + matern_52_kernel, + 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): +# TODO: deprecate +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_orthonormality(backend): + eigenfunctions = Circle().get_eigenfunctions(10) 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 + check_function_with_backend( + backend, + np.eye(eigenfunctions.num_eigenfunctions), + lambda theta: B.matmul(B.T(eigenfunctions(theta)), eigenfunctions(theta)) + / len(theta), + theta, + atol=1e-2, ) - 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)) +@pytest.mark.parametrize("nu, atol", [(0.5, 1), (1.5, 3), (2.5, 5), (np.inf, 6)]) +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_equivalence_kernel(nu, atol, backend): - -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) if nu == 0.5: - return B.exp(-r) + analytic_kernel = matern_12_kernel elif nu == 1.5: - sqrt3 = np.sqrt(3.0) - return (1.0 + sqrt3 * r) * B.exp(-sqrt3 * r) + analytic_kernel = 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 = matern_52_kernel elif nu == np.inf: - return B.exp(-0.5 * r**2) - else: - raise NotImplementedError - + analytic_kernel = rbf_kernel -@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]) + inputs = np.random.uniform(0, 2 * np.pi, size=(5, 1)) + inputs2 = np.random.uniform(0, 2 * np.pi, size=(3, 1)) - 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, + lambda nu, lengthscale, X, X2: kernel.K( + {"nu": nu, "lengthscale": lengthscale}, X, X2 + ), + np.array([nu]), + np.array([1.0]), + inputs, + inputs2, + atol=atol, ) diff --git a/tests/spaces/test_eigenfunctions_basics.py b/tests/spaces/test_eigenfunctions_basics.py new file mode 100644 index 00000000..f86a3f0d --- /dev/null +++ b/tests/spaces/test_eigenfunctions_basics.py @@ -0,0 +1,220 @@ +import jax +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.matern_kernel import default_num +from geometric_kernels.spaces import ( + Circle, + CompactMatrixLieGroup, + HypercubeGraph, + Hypersphere, + SpecialOrthogonal, + SpecialUnitary, +) +from geometric_kernels.utils.utils import chain + +from ..helper import check_function_with_backend + +jax.config.update("jax_enable_x64", True) # enable float64 in JAX + + +@pytest.fixture( + params=[ + Circle(), + HypercubeGraph(1), + HypercubeGraph(3), + HypercubeGraph(10), + Hypersphere(2), + Hypersphere(3), + Hypersphere(10), + SpecialOrthogonal(3), + SpecialOrthogonal(8), + SpecialUnitary(2), + SpecialUnitary(5), + ], + ids=str, +) +def inputs(request) -> Tuple[B.Numeric]: + """ + 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, + - 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 + ) + eigenfunctions = space.get_eigenfunctions(num_levels) + + key = np.random.RandomState() + N, N2 = key.randint(low=1, high=100 + 1, size=2) + # N, N2 = 2, 3 # TODO: fixme + 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 + 0.01 + + 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() + 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"]) +@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_hypercube_graph.py b/tests/spaces/test_hypercube_graph.py index 761fe614..92c441a1 100644 --- a/tests/spaces/test_hypercube_graph.py +++ b/tests/spaces/test_hypercube_graph.py @@ -1,13 +1,12 @@ 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.utils import binary_vectors_and_subsets from ..helper import check_function_with_backend @@ -38,42 +37,17 @@ def inputs(request) -> Tuple[B.Numeric]: return space, eigenfunctions, X, X2, weights -@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, - ) - - def test_numbers_of_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 - +# TODO: deprecate @pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) def test_orthonormality(inputs, backend): space, _, _, _, _ = inputs @@ -95,93 +69,6 @@ def test_orthonormality(inputs, backend): ) -@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) -def test_weighted_outerproduct_with_addition_theorem(inputs, backend): - _, eigenfunctions, X, X2, weights = inputs - - 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, _, 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, - ) - - -@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 = 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, _, weights = inputs - - phi_product_diag = eigenfunctions.phi_product_diag(X) - - 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): diff --git a/tests/spaces/test_hypersphere.py b/tests/spaces/test_hypersphere.py index 61f51bfd..9372e5bb 100644 --- a/tests/spaces/test_hypersphere.py +++ b/tests/spaces/test_hypersphere.py @@ -1,122 +1,37 @@ import lab as B import numpy as np import pytest -from plum import Tuple from geometric_kernels.spaces import Hypersphere -from geometric_kernels.spaces.hypersphere import SphericalHarmonics -from geometric_kernels.utils.utils import chain - - -class Consts: - seed = 42 - dimension = 2 - num_data = 7 - num_data2 = 5 - num_eigenfunctions = 9 - - -@pytest.fixture(name="eigenfunctions") -def _eigenfunctions_fixture(): - return SphericalHarmonics(Consts.dimension, Consts.num_eigenfunctions) - - -@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 - - 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 - ) - 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)) +from ..helper import check_function_with_backend + + +# TODO: deprecate +@pytest.mark.parametrize("dim, num_levels", [(2, 5), (3, 5), (19, 2)]) +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_orthonormality(dim, num_levels, backend): + space = Hypersphere(dim) + eigenfunctions = space.get_eigenfunctions(num_levels) + + key = np.random.RandomState(1234) + + key, xs = space.random(key, 10000) + + # print(B.matmul(B.T(eigenfunctions(xs)), eigenfunctions(xs)) / xs.shape[0]) + + 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=1e-1, + ) 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 738d1f63fafe4179c78bbe3beaae467a979f6aa6 Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Mon, 21 Oct 2024 16:24:26 +0200 Subject: [PATCH 06/35] Remove deprecated tests --- tests/spaces/test_circle.py | 17 ----------------- tests/spaces/test_hypercube_graph.py | 23 ----------------------- tests/spaces/test_hypersphere.py | 25 ------------------------- 3 files changed, 65 deletions(-) diff --git a/tests/spaces/test_circle.py b/tests/spaces/test_circle.py index 665bdebf..f6e6689c 100644 --- a/tests/spaces/test_circle.py +++ b/tests/spaces/test_circle.py @@ -14,26 +14,9 @@ from ..helper import check_function_with_backend -# TODO: deprecate -@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) -def test_orthonormality(backend): - eigenfunctions = Circle().get_eigenfunctions(10) - theta = np.linspace(-np.pi, np.pi, 5_000).reshape(-1, 1) # [N, 1] - - check_function_with_backend( - backend, - np.eye(eigenfunctions.num_eigenfunctions), - lambda theta: B.matmul(B.T(eigenfunctions(theta)), eigenfunctions(theta)) - / len(theta), - theta, - atol=1e-2, - ) - - @pytest.mark.parametrize("nu, atol", [(0.5, 1), (1.5, 3), (2.5, 5), (np.inf, 6)]) @pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) def test_equivalence_kernel(nu, atol, backend): - if nu == 0.5: analytic_kernel = matern_12_kernel elif nu == 1.5: diff --git a/tests/spaces/test_hypercube_graph.py b/tests/spaces/test_hypercube_graph.py index 92c441a1..73296894 100644 --- a/tests/spaces/test_hypercube_graph.py +++ b/tests/spaces/test_hypercube_graph.py @@ -6,7 +6,6 @@ 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 from ..helper import check_function_with_backend @@ -47,28 +46,6 @@ def test_numbers_of_eigenfunctions(inputs): assert eigenfunctions.num_eigenfunctions == 2**space.dim -# TODO: deprecate -@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("lengthscale", [1.0, 5.0, 10.0]) @pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) def test_against_analytic_heat_kernel(inputs, lengthscale, backend): diff --git a/tests/spaces/test_hypersphere.py b/tests/spaces/test_hypersphere.py index 9372e5bb..7cfb20fe 100644 --- a/tests/spaces/test_hypersphere.py +++ b/tests/spaces/test_hypersphere.py @@ -1,33 +1,8 @@ import lab as B import numpy as np -import pytest from geometric_kernels.spaces import Hypersphere -from ..helper import check_function_with_backend - - -# TODO: deprecate -@pytest.mark.parametrize("dim, num_levels", [(2, 5), (3, 5), (19, 2)]) -@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) -def test_orthonormality(dim, num_levels, backend): - space = Hypersphere(dim) - eigenfunctions = space.get_eigenfunctions(num_levels) - - key = np.random.RandomState(1234) - - key, xs = space.random(key, 10000) - - # print(B.matmul(B.T(eigenfunctions(xs)), eigenfunctions(xs)) / xs.shape[0]) - - 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=1e-1, - ) - def test_sphere_heat_kernel(): # Tests that the heat kernel on the sphere solves the heat equation. This From 575b11b8595220462938137ca643f10a2d423388 Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Mon, 21 Oct 2024 22:00:56 +0200 Subject: [PATCH 07/35] Add Graph and Mesh to tests/spaces/test_eigenfunctions_basics.py --- tests/helper.py | 31 ++++++++++++++++++++++ tests/spaces/test_eigenfunctions_basics.py | 13 +++++++-- tests/spaces/test_graph.py | 27 +++---------------- tests/spaces/test_mesh.py | 14 +++------- 4 files changed, 50 insertions(+), 35 deletions(-) diff --git a/tests/helper.py b/tests/helper.py index 3dad2210..0c877643 100644 --- a/tests/helper.py +++ b/tests/helper.py @@ -1,3 +1,5 @@ +from pathlib import Path + import lab as B import numpy as np from beartype.door import die_if_unbearable, is_bearable @@ -6,6 +8,35 @@ EagerTensor = ModuleType("tensorflow.python.framework.ops", "EagerTensor") +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 + def np_to_backend(value: B.NPNumeric, backend: str): """ diff --git a/tests/spaces/test_eigenfunctions_basics.py b/tests/spaces/test_eigenfunctions_basics.py index f86a3f0d..c998af7b 100644 --- a/tests/spaces/test_eigenfunctions_basics.py +++ b/tests/spaces/test_eigenfunctions_basics.py @@ -9,14 +9,16 @@ from geometric_kernels.spaces import ( Circle, CompactMatrixLieGroup, + Graph, HypercubeGraph, Hypersphere, + Mesh, SpecialOrthogonal, SpecialUnitary, ) from geometric_kernels.utils.utils import chain -from ..helper import check_function_with_backend +from ..helper import TEST_GRAPH_ADJACENCY, TEST_MESH_PATH, check_function_with_backend jax.config.update("jax_enable_x64", True) # enable float64 in JAX @@ -26,7 +28,7 @@ Circle(), HypercubeGraph(1), HypercubeGraph(3), - HypercubeGraph(10), + HypercubeGraph(6), Hypersphere(2), Hypersphere(3), Hypersphere(10), @@ -34,6 +36,9 @@ SpecialOrthogonal(8), SpecialUnitary(2), SpecialUnitary(5), + Mesh.load_mesh(TEST_MESH_PATH), + Graph(TEST_GRAPH_ADJACENCY, normalize_laplacian=False), + Graph(TEST_GRAPH_ADJACENCY, normalize_laplacian=True), ], ids=str, ) @@ -55,6 +60,10 @@ def inputs(request) -> Tuple[B.Numeric]: 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() diff --git a/tests/spaces/test_graph.py b/tests/spaces/test_graph.py index fc94155c..dd4a097e 100644 --- a/tests/spaces/test_graph.py +++ b/tests/spaces/test_graph.py @@ -13,31 +13,12 @@ from geometric_kernels.tensorflow import * # noqa from geometric_kernels.torch import * # noqa +from ..helper import TEST_GRAPH_ADJACENCY, TEST_GRAPH_LAPLACIAN + 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) +A = TEST_GRAPH_ADJACENCY +L = TEST_GRAPH_LAPLACIAN def run_tests_with_adj(A, L, tol=1e-7, tol_m=1e-4): 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)) From e952af897863d1902f0e8c04e6cecafcf6efe344 Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Thu, 24 Oct 2024 10:51:14 +0200 Subject: [PATCH 08/35] Fix ProductDiscreteSpectrumSpace's incompatibility with the TensorFlow backend --- geometric_kernels/spaces/product.py | 43 ++++++++++++++--------------- 1 file changed, 20 insertions(+), 23 deletions(-) diff --git a/geometric_kernels/spaces/product.py b/geometric_kernels/spaces/product.py index 0d99920b..74d2b3ee 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 @@ -286,20 +286,19 @@ def phi_product( phis = B.stack( *[ - eigenfunction.phi_product(X1, X2, **kwargs) - for eigenfunction, X1, X2 in zip(self.eigenfunctions, Xs, Xs2) + take_along_axis( + eigenfunction.phi_product(X1, X2, **kwargs), + from_numpy(X1, self.eigenindicies[None, None, :, s]), + -1, + ) + for s, (eigenfunction, X1, X2) in enumerate( + zip(self.eigenfunctions, Xs, Xs2) + ) ], axis=-1, - ) # [N, N2, LFactor, S] where `LFactor` is self.eigenfunctions[0].num_levels + ) # [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 @@ -308,19 +307,17 @@ def phi_product_diag(self, X: B.Numeric, **kwargs): 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), + from_numpy(X1, self.eigenindicies[None, :, s]), + -1, + ) + for s, (eigenfunction, X1) in enumerate(zip(self.eigenfunctions, Xs)) ], axis=-1, - ) # [N, LFactor, S] where `LFactor` is self.eigenfunctions[0].num_levels - - prod_phis = phis[ - :, - self.eigenindicies, - B.range(self.eigenindicies.shape[1]), - ].prod( - axis=-1 - ) # [N, LFactor, S] -> [N, L] + ) # [N, L, S] + + prod_phis = B.prod(phis, axis=-1) # [N, L, S] -> [N, L] return prod_phis From 9c5b65e75878b988060601503ae422e10a0e2944 Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Thu, 24 Oct 2024 10:57:45 +0200 Subject: [PATCH 09/35] Graph now remembers if the Laplacian is normalized and complaints if the requested number of eigenpairs exceeds the size of the adjacency matrix. Also, added support of sparray (in addition to spmatrix) in lab_extras/numpy/sparse_extras.py. --- geometric_kernels/lab_extras/numpy/sparse_extras.py | 9 ++------- geometric_kernels/spaces/graph.py | 6 +++++- 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/geometric_kernels/lab_extras/numpy/sparse_extras.py b/geometric_kernels/lab_extras/numpy/sparse_extras.py index 19c316be..f8c9e8c6 100644 --- a/geometric_kernels/lab_extras/numpy/sparse_extras.py +++ b/geometric_kernels/lab_extras/numpy/sparse_extras.py @@ -12,13 +12,8 @@ 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, + sp.sparray, + sp.spmatrix, ] diff --git a/geometric_kernels/spaces/graph.py b/geometric_kernels/spaces/graph.py index bbf92f4d..d93ab7ac 100644 --- a/geometric_kernels/spaces/graph.py +++ b/geometric_kernels/spaces/graph.py @@ -60,9 +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})" + return f"Graph({self.num_vertices}, {'normalized' if self._normalized else 'unnormalized'})" @staticmethod def _checks(adjacency): @@ -118,6 +119,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) From 77596fac6db76a7045378a829e19bed91d889376 Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Thu, 24 Oct 2024 11:18:14 +0200 Subject: [PATCH 10/35] Revise tests/spaces/test_graph.py and tests/spaces/test_eigenfunctions_basics.py, add tests/spaces/test_eigenvalues.py and tests/spaces/test_basics.py --- tests/data.py | 47 ++++ tests/helper.py | 84 ++++--- tests/spaces/test_basics.py | 28 +++ tests/spaces/test_eigenfunctions_basics.py | 37 +-- tests/spaces/test_eigenvalues.py | 53 +++++ tests/spaces/test_graph.py | 253 +++++++-------------- 6 files changed, 272 insertions(+), 230 deletions(-) create mode 100644 tests/data.py create mode 100644 tests/spaces/test_basics.py create mode 100644 tests/spaces/test_eigenvalues.py 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/helper.py b/tests/helper.py index 0c877643..e4610601 100644 --- a/tests/helper.py +++ b/tests/helper.py @@ -1,41 +1,49 @@ -from pathlib import Path - 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, + DiscreteSpectrumSpace, + Graph, + HypercubeGraph, + Hypersphere, + Mesh, + ProductDiscreteSpectrumSpace, + SpecialOrthogonal, + SpecialUnitary, +) + +from .data import TEST_GRAPH_ADJACENCY, TEST_MESH_PATH + EagerTensor = ModuleType("tensorflow.python.framework.ops", "EagerTensor") -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], + +def discrete_spectrum_spaces() -> List[DiscreteSpectrumSpace]: + return [ + Circle(), + HypercubeGraph(1), + HypercubeGraph(3), + HypercubeGraph(6), + Hypersphere(2), + Hypersphere(3), + Hypersphere(10), + SpecialOrthogonal(3), + SpecialOrthogonal(8), + SpecialUnitary(2), + SpecialUnitary(5), + Mesh.load_mesh(TEST_MESH_PATH), + Graph(TEST_GRAPH_ADJACENCY, normalize_laplacian=False), + Graph(TEST_GRAPH_ADJACENCY, normalize_laplacian=True), + 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)), ] -).astype( - np.float64 -) # corresponds to TEST_GRAPH_ADJACENCY, unnormalized Laplacian def np_to_backend(value: B.NPNumeric, backend: str): @@ -67,6 +75,10 @@ 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)) @@ -76,7 +88,8 @@ 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. @@ -89,6 +102,8 @@ 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}") @@ -139,6 +154,11 @@ def check_function_with_backend( f_output = f(*args_casted) assert is_bearable(f_output, array_type(backend)) 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) + np.testing.assert_allclose(f_output, result, atol=atol) else: assert compare_to_result(result, f_output) diff --git a/tests/spaces/test_basics.py b/tests/spaces/test_basics.py new file mode 100644 index 00000000..981447e9 --- /dev/null +++ b/tests/spaces/test_basics.py @@ -0,0 +1,28 @@ +import inspect + +import geometric_kernels.spaces + +from ..helper import discrete_spectrum_spaces + + +def test_all_discrete_spectrum_spaces_covered(): + spaces = discrete_spectrum_spaces() + + # 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, geometric_kernels.spaces.DiscreteSpectrumSpace + ) 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"Space {cls_name} not covered by tests" diff --git a/tests/spaces/test_eigenfunctions_basics.py b/tests/spaces/test_eigenfunctions_basics.py index c998af7b..4fc21053 100644 --- a/tests/spaces/test_eigenfunctions_basics.py +++ b/tests/spaces/test_eigenfunctions_basics.py @@ -3,48 +3,23 @@ import numpy as np import pytest from opt_einsum import contract as einsum -from plum import Tuple from geometric_kernels.kernels.matern_kernel import default_num -from geometric_kernels.spaces import ( - Circle, - CompactMatrixLieGroup, - Graph, - HypercubeGraph, - Hypersphere, - Mesh, - SpecialOrthogonal, - SpecialUnitary, -) +from geometric_kernels.spaces import CompactMatrixLieGroup, Hypersphere, Mesh from geometric_kernels.utils.utils import chain -from ..helper import TEST_GRAPH_ADJACENCY, TEST_MESH_PATH, check_function_with_backend +from ..helper import check_function_with_backend, discrete_spectrum_spaces jax.config.update("jax_enable_x64", True) # enable float64 in JAX @pytest.fixture( - params=[ - Circle(), - HypercubeGraph(1), - HypercubeGraph(3), - HypercubeGraph(6), - Hypersphere(2), - Hypersphere(3), - Hypersphere(10), - SpecialOrthogonal(3), - SpecialOrthogonal(8), - SpecialUnitary(2), - SpecialUnitary(5), - Mesh.load_mesh(TEST_MESH_PATH), - Graph(TEST_GRAPH_ADJACENCY, normalize_laplacian=False), - Graph(TEST_GRAPH_ADJACENCY, normalize_laplacian=True), - ], + params=discrete_spectrum_spaces(), ids=str, ) -def inputs(request) -> Tuple[B.Numeric]: +def inputs(request): """ - Returns a tuple (space, eigenfunctions, X, X2) where: + 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, @@ -68,7 +43,6 @@ def inputs(request) -> Tuple[B.Numeric]: key = np.random.RandomState() N, N2 = key.randint(low=1, high=100 + 1, size=2) - # N, N2 = 2, 3 # TODO: fixme key, X = space.random(key, N) key, X2 = space.random(key, N2) @@ -164,7 +138,6 @@ def test_weighted_outerproduct_with_addition_theorem(inputs, backend): ) -# @pytest.mark.parametrize("backend", ["numpy"]) @pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) def test_weighted_outerproduct_with_addition_theorem_one_input(inputs, backend): _, eigenfunctions, X, _, weights = inputs diff --git a/tests/spaces/test_eigenvalues.py b/tests/spaces/test_eigenvalues.py new file mode 100644 index 00000000..bccc02f2 --- /dev/null +++ b/tests/spaces/test_eigenvalues.py @@ -0,0 +1,53 @@ +""" +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). + """ + 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 dd4a097e..a31fa42a 100644 --- a/tests/spaces/test_graph.py +++ b/tests/spaces/test_graph.py @@ -1,19 +1,21 @@ 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 -from ..helper import TEST_GRAPH_ADJACENCY, TEST_GRAPH_LAPLACIAN +from ..data import ( + TEST_GRAPH_ADJACENCY, + TEST_GRAPH_LAPLACIAN, + TEST_GRAPH_LAPLACIAN_NORMALIZED, +) +from ..helper import check_function_with_backend, np_to_backend warnings.filterwarnings("ignore", category=RuntimeWarning, module="scipy") @@ -21,181 +23,100 @@ L = TEST_GRAPH_LAPLACIAN -def run_tests_with_adj(A, L, tol=1e-7, tol_m=1e-4): - ############################################## - # Inits +@pytest.mark.parametrize("normalized", [True, False]) +@pytest.mark.parametrize( + "backend", ["numpy", "tensorflow", "torch", "jax", "scipy_sparse"] +) +def test_laplacian(normalized, backend): - n = A.shape[0] - L = B.cast(B.dtype(A), L) - graph = Graph(A) + # 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, + ) - normed_graph = Graph(A, normalize_laplacian=True) - ############################################## - # Laplacian computation +@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, + ) - comparison = graph._laplacian == L + def eigendiff(adj): + graph = Graph(adj, normalize_laplacian=normalized) - if sp.issparse(comparison): - comparison = comparison.toarray() + 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 - if isinstance(comparison, np.matrix): # bug with lab? - assert comparison.all(), "Laplacian does not match." - else: - assert B.all(comparison), "Laplacian does not match." + eigenvalue_mat = sp.csr_array(eigenvalue_mat) + eigenvectors = sp.csr_array(eigenvectors) - normed_l = normed_graph._laplacian - if sp.issparse(normed_l): - normed_l = normed_l.toarray() + laplace_x_eigvecs = laplacian @ eigenvectors + eigvals_x_eigvecs = eigenvectors @ eigenvalue_mat + return laplace_x_eigvecs - eigvals_x_eigvecs - assert ( - B.max(B.abs(B.diag(normed_l) - 1)[:-1]) < tol_m - and B.abs(B.diag(normed_l)[-1] - 0) < tol_m + check_function_with_backend( + backend, + np.zeros((TEST_GRAPH_ADJACENCY.shape[0], L)), + eigendiff, + TEST_GRAPH_ADJACENCY, ) - ############################################## - # Eigendecomposition checks - evecs = graph.get_eigenvectors(n) - evals = graph.get_eigenvalues(n) +@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): - 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) - - # check vecs - np.testing.assert_allclose( - np.abs(evecs)[:, 2:], np.abs(evecs_np)[:, 2:], atol=tol, rtol=tol + 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, rtol=tol - ) - except AssertionError: - np.testing.assert_allclose( - np.abs(evecs)[:, [1, 0]], np.abs(evecs_np)[:, :2], atol=tol, rtol=tol + 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], ) - 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 - - 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) - - normed_params["nu"], normed_params["lengthscale"] = nu, lscale - K_normed_cons.K(normed_params, idx) - - ############################################## - # Matern 2 check - - 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) - - ############################################## - # RBF check - - 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 - - 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) - - try: - np.testing.assert_allclose( - np.abs(evecs)[:, :2], np.abs(evecs_np)[:, :2], atol=tol_m, rtol=tol_m - ) - except AssertionError: - np.testing.assert_allclose( - np.abs(evecs)[:, [1, 0]], np.abs(evecs_np)[:, :2], atol=tol_m, rtol=tol_m + if nu < np.inf: + K = ( + evecs_np + @ np.diag(np.power(evals_np + 2 * nu / lengthscale**2, -nu)) + @ evecs_np.T ) - - 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) 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]), + ) From 1d368921f0b1bcfad4892e8768039b160f7f8b66 Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Thu, 24 Oct 2024 11:20:33 +0200 Subject: [PATCH 11/35] Stop tests/kernels/test_normalization.py from requesting more than the maximal number of eigenpairs for a graph --- tests/kernels/test_normalization.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/kernels/test_normalization.py b/tests/kernels/test_normalization.py index c95d02a3..357748f4 100644 --- a/tests/kernels/test_normalization.py +++ b/tests/kernels/test_normalization.py @@ -44,6 +44,7 @@ def test_normalization_matern_kl_kernel(space_name): space = Graph(A, normalize_laplacian=True) points = np.arange(space.num_vertices).reshape(-1, 1) + num_eigenfns = min(num_eigenfns, A.shape[0]) else: raise ValueError(f"Unknown space {space}") From 50caf1bfd7ff8645c7177877f7cfe1dfccec7cd0 Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Thu, 24 Oct 2024 11:22:31 +0200 Subject: [PATCH 12/35] Rename test_eigenfunctions_basics.py to test_eigenfunctions.py --- .../{test_eigenfunctions_basics.py => test_eigenfunctions.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tests/spaces/{test_eigenfunctions_basics.py => test_eigenfunctions.py} (100%) diff --git a/tests/spaces/test_eigenfunctions_basics.py b/tests/spaces/test_eigenfunctions.py similarity index 100% rename from tests/spaces/test_eigenfunctions_basics.py rename to tests/spaces/test_eigenfunctions.py From 3b6799ea496ebf04711e2e6daadcfb1b95e5fcae Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Thu, 24 Oct 2024 12:09:28 +0200 Subject: [PATCH 13/35] Python 3.8 compatibility --- .../lab_extras/numpy/sparse_extras.py | 21 +++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/geometric_kernels/lab_extras/numpy/sparse_extras.py b/geometric_kernels/lab_extras/numpy/sparse_extras.py index f8c9e8c6..f8be84ad 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,10 +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.sparray, - sp.spmatrix, -] +if sys.version_info[:2] <= (3, 9): + 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 From 9281d5e548c8c459206fd0d5fa13dc9c5f110d79 Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Fri, 25 Oct 2024 16:40:46 +0200 Subject: [PATCH 14/35] Fix Graph checks for the scipy_sparse backend --- geometric_kernels/lab_extras/numpy/sparse_extras.py | 2 +- geometric_kernels/spaces/graph.py | 5 ++--- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/geometric_kernels/lab_extras/numpy/sparse_extras.py b/geometric_kernels/lab_extras/numpy/sparse_extras.py index f8be84ad..e4a938ce 100644 --- a/geometric_kernels/lab_extras/numpy/sparse_extras.py +++ b/geometric_kernels/lab_extras/numpy/sparse_extras.py @@ -13,7 +13,7 @@ SparseArray defines a lab data type that covers all possible sparse scipy arrays, so that multiple dispatch works with such arrays. """ -if sys.version_info[:2] <= (3, 9): +if sys.version_info[:2] <= (3, 8): SparseArray = Union[ sp.bsr_matrix, sp.coo_matrix, diff --git a/geometric_kernels/spaces/graph.py b/geometric_kernels/spaces/graph.py index d93ab7ac..94cf71da 100644 --- a/geometric_kernels/spaces/graph.py +++ b/geometric_kernels/spaces/graph.py @@ -71,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: From d392bc099b302d600f8cf0f50164d2fa7e15c64b Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Sun, 27 Oct 2024 16:58:53 +0100 Subject: [PATCH 15/35] Move MaternKarhunenLoeveKernel import inside DeterministicFeatureMapCompact's __init__ to avoid circular imports --- geometric_kernels/feature_maps/deterministic.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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) From a817021dcda588898a951138b73081b6e105437f Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Sun, 27 Oct 2024 23:27:45 +0100 Subject: [PATCH 16/35] Revise test_lie_groups.py. test_characters_orthogonal currently fails because of issue #152 --- tests/helper.py | 37 ++++-- tests/spaces/test_lie_groups.py | 203 ++++++++++++++++---------------- 2 files changed, 129 insertions(+), 111 deletions(-) diff --git a/tests/helper.py b/tests/helper.py index e4610601..7823ca36 100644 --- a/tests/helper.py +++ b/tests/helper.py @@ -7,6 +7,7 @@ from geometric_kernels.lab_extras import SparseArray from geometric_kernels.spaces import ( Circle, + CompactMatrixLieGroup, DiscreteSpectrumSpace, Graph, HypercubeGraph, @@ -22,22 +23,17 @@ EagerTensor = ModuleType("tensorflow.python.framework.ops", "EagerTensor") -def discrete_spectrum_spaces() -> List[DiscreteSpectrumSpace]: +def compact_matrix_lie_groups() -> List[CompactMatrixLieGroup]: return [ - Circle(), - HypercubeGraph(1), - HypercubeGraph(3), - HypercubeGraph(6), - Hypersphere(2), - Hypersphere(3), - Hypersphere(10), SpecialOrthogonal(3), SpecialOrthogonal(8), SpecialUnitary(2), SpecialUnitary(5), - Mesh.load_mesh(TEST_MESH_PATH), - Graph(TEST_GRAPH_ADJACENCY, normalize_laplacian=False), - Graph(TEST_GRAPH_ADJACENCY, normalize_laplacian=True), + ] + + +def product_discrete_spectrum_spaces() -> List[ProductDiscreteSpectrumSpace]: + return [ ProductDiscreteSpectrumSpace(Circle(), Hypersphere(3), Circle()), ProductDiscreteSpectrumSpace( Circle(), Graph(np.kron(TEST_GRAPH_ADJACENCY, TEST_GRAPH_ADJACENCY)) @@ -46,6 +42,25 @@ def discrete_spectrum_spaces() -> List[DiscreteSpectrumSpace]: ] +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 np_to_backend(value: B.NPNumeric, backend: str): """ Converts a numpy array to the desired backend. diff --git a/tests/spaces/test_lie_groups.py b/tests/spaces/test_lie_groups.py index 863580d8..d8ff3da3 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(5, default_num(space)) + eigenfunctions = space.get_eigenfunctions(num_levels) + key = np.random.RandomState() + 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,93 @@ 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 + + # if isinstance(group, SpecialOrthogonal) and group.n % 2 == 0: + # pytest.skip("This test currently fails for SO(n) with n even, because of a bug") + + num_samples = 10000 + key = np.random.RandomState() + _, 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 + ) From 461670378ea254a7e286388fa932ba640f6949de Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Mon, 28 Oct 2024 20:36:15 +0100 Subject: [PATCH 17/35] Fix a bug in SpecialOrthogonal introduced by cc85678. Fixes issue #152 --- geometric_kernels/spaces/so.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/geometric_kernels/spaces/so.py b/geometric_kernels/spaces/so.py index 90c3c945..847b399a 100644 --- a/geometric_kernels/spaces/so.py +++ b/geometric_kernels/spaces/so.py @@ -166,7 +166,7 @@ def _torus_representative(self, X: B.Numeric) -> B.Numeric: ) # eigenvectors calculated by LAPACK are either real or purely imaginary, make everything real # WARNING: might depend on the implementation of the eigendecomposition! - c = B.real(c) + B.real(c) + c = B.real(c) + B.imag(c) # normalize s.t. det(c)≈±1, probably unnecessary c /= math.sqrt(2) eigvals = B.concat( From 94b29b10aa783ff542ee0d3a67e3a8c5bbb05e68 Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Mon, 28 Oct 2024 20:37:02 +0100 Subject: [PATCH 18/35] Finish revising tests/spaces/test_lie_groups.py --- tests/spaces/test_lie_groups.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/tests/spaces/test_lie_groups.py b/tests/spaces/test_lie_groups.py index d8ff3da3..f64b5a41 100644 --- a/tests/spaces/test_lie_groups.py +++ b/tests/spaces/test_lie_groups.py @@ -22,7 +22,7 @@ def inputs(request): - X2 is another random sample of random size from the space, """ space = request.param - num_levels = min(5, default_num(space)) + num_levels = min(10, default_num(space)) eigenfunctions = space.get_eigenfunctions(num_levels) key = np.random.RandomState() @@ -109,9 +109,6 @@ def test_character_at_identity(inputs, backend): def test_characters_orthogonal(inputs, backend): group, eigenfunctions, _, _ = inputs - # if isinstance(group, SpecialOrthogonal) and group.n % 2 == 0: - # pytest.skip("This test currently fails for SO(n) with n even, because of a bug") - num_samples = 10000 key = np.random.RandomState() _, X = group.random(key, num_samples) From c95842e44f3af151204de4adca18495337b09f86 Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Mon, 11 Nov 2024 16:34:44 +0100 Subject: [PATCH 19/35] Revise tests/kernels/test_matern_karhunenloeve_kernel.py. Make check_function_with_backend with custom compare_to_result more verbose --- tests/helper.py | 4 +- .../test_matern_karhunenloeve_kernel.py | 213 +++++++++++++++--- 2 files changed, 181 insertions(+), 36 deletions(-) diff --git a/tests/helper.py b/tests/helper.py index 7823ca36..5fadfa63 100644 --- a/tests/helper.py +++ b/tests/helper.py @@ -176,4 +176,6 @@ def check_function_with_backend( f_output = B.to_numpy(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_matern_karhunenloeve_kernel.py b/tests/kernels/test_matern_karhunenloeve_kernel.py index 1d87ecb2..662fb040 100644 --- a/tests/kernels/test_matern_karhunenloeve_kernel.py +++ b/tests/kernels/test_matern_karhunenloeve_kernel.py @@ -1,62 +1,205 @@ -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, + - num_levels = default_num(space), + - kernel = MaternKarhunenLoeveKernel(space, 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, 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() + 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), + lambda nu, lengthscale: kernel.eigenvalues( + {"nu": nu, "lengthscale": lengthscale} + ), + params["nu"], + params["lengthscale"], + compare_to_result=lambda res, f_out: f_out.shape == 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, + lambda nu, lengthscale: kernel.eigenvalues( + {"nu": nu, "lengthscale": lengthscale} + ), + params["nu"], + params["lengthscale"], + compare_to_result=lambda _, f_out: np.all(B.to_numpy(f_out) >= 0), + ) + + +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_eigenvalues_ordered(inputs, backend): + _, _, kernel, _, _ = inputs + params = kernel.init_params() -def test_K_shapes(kernel: MaternKarhunenLoeveKernel): + # Check that the eigenvalues are sorted in descending order. + check_function_with_backend( + backend, + None, + lambda nu, lengthscale: kernel.eigenvalues( + {"nu": nu, "lengthscale": lengthscale} + ), + params["nu"], + params["lengthscale"], + 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) + + assert result.shape == (X.shape[0], X2.shape[0]), "K has incorrect shape" - K = kernel.K(params, X, None) - assert K.shape == (N1, N1) + if backend != "numpy": + # Check that kernel.K computed using `backend` coincides with the numpy result. + check_function_with_backend( + backend, + result, + lambda nu, lengthscale, X, X2: kernel.K( + {"nu": nu, "lengthscale": lengthscale}, X, X2 + ), + params["nu"], + params["lengthscale"], + X, + X2, + ) - K = kernel.K(params, X, X2) - assert K.shape == (N1, N2) - K = kernel.K_diag(params, X) - assert K.shape == (N1,) +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_K_one_param(inputs, backend): + space, _, kernel, X, _ = inputs + params = kernel.init_params() + result = kernel.K(params, X, X) -def test_normalize(kernel: MaternKarhunenLoeveKernel): - random_points = np.arange(kernel.space.num_vertices).reshape(-1, 1) + # Check that kernel.K(X) coincides with kernel.K(X, X). + check_function_with_backend( + backend, + result, + lambda nu, lengthscale, X: kernel.K({"nu": nu, "lengthscale": lengthscale}, X), + params["nu"], + params["lengthscale"], + X, + atol=1e-2 if isinstance(space, Mesh) else _EPS, + ) + +@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() + + assert result.shape == (X.shape[0],), "The diagonal has incorrect shape" - np.testing.assert_allclose(np.mean(K), 1.0) + # Check that kernel.K_diag coincides with the diagonal of kernel.K. + check_function_with_backend( + backend, + result, + lambda nu, lengthscale, X: kernel.K_diag( + {"nu": nu, "lengthscale": lengthscale}, X + ), + params["nu"], + params["lengthscale"], + 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() + key, X = space.random( + key, 1000 + ) # we need a large sample to get a good estimate of the mean variance + + # Check that the average variance of the kernel is 1. + check_function_with_backend( + backend, + np.array(1.0), + lambda nu, lengthscale, X: B.mean( + kernel.K_diag({"nu": nu, "lengthscale": lengthscale}, X), squeeze=False + ), + params["nu"], + params["lengthscale"], + X, + atol=0.1, # very loose, but helps make sure the result is close to 1 + ) From f86e34bef2b00ea03991f5f79ed94a456410daba Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Mon, 11 Nov 2024 17:01:23 +0100 Subject: [PATCH 20/35] Minor fixes in tests/kernels/test_matern_karhunenloeve_kernel.py --- tests/kernels/test_matern_karhunenloeve_kernel.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/kernels/test_matern_karhunenloeve_kernel.py b/tests/kernels/test_matern_karhunenloeve_kernel.py index 662fb040..fcee3cf5 100644 --- a/tests/kernels/test_matern_karhunenloeve_kernel.py +++ b/tests/kernels/test_matern_karhunenloeve_kernel.py @@ -21,9 +21,9 @@ def inputs(request): """ Returns a tuple (space, num_levels, kernel, X, X2) where: - - space = request.param, + - space = request.param[0], - num_levels = default_num(space), - - kernel = MaternKarhunenLoeveKernel(space, num_levels), + - 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, """ @@ -201,5 +201,5 @@ def test_normalize(inputs, backend): params["nu"], params["lengthscale"], X, - atol=0.1, # very loose, but helps make sure the result is close to 1 + atol=0.2, # very loose, but helps make sure the result is close to 1 ) From 9b2c3184b3e5d2230f5fd89632d63c6ff936b2f5 Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Mon, 11 Nov 2024 18:35:23 +0100 Subject: [PATCH 21/35] Fix __str__ in geometric_kernels/spaces/spd.py --- geometric_kernels/spaces/spd.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/geometric_kernels/spaces/spd.py b/geometric_kernels/spaces/spd.py index 7cd9bb34..3135d98d 100644 --- a/geometric_kernels/spaces/spd.py +++ b/geometric_kernels/spaces/spd.py @@ -55,7 +55,7 @@ def __init__(self, n): super().__init__(n) def __str__(self): - return f"SpecialOrthogonal({self.n})" + return f"SymmetricPositiveDefiniteMatrices({self.n})" @property def dimension(self) -> int: From 45157f128afd475227f0c74ab69a85b75888b121 Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Mon, 11 Nov 2024 18:36:32 +0100 Subject: [PATCH 22/35] New tests for MaternFeatureMapKernel: tests/kernels/test_feature_map_kernel.py --- tests/helper.py | 21 ++++ tests/kernels/test_feature_map_kernel.py | 135 +++++++++++++++++++++++ 2 files changed, 156 insertions(+) create mode 100644 tests/kernels/test_feature_map_kernel.py diff --git a/tests/helper.py b/tests/helper.py index 5fadfa63..30ae4818 100644 --- a/tests/helper.py +++ b/tests/helper.py @@ -10,12 +10,15 @@ CompactMatrixLieGroup, DiscreteSpectrumSpace, Graph, + Hyperbolic, HypercubeGraph, Hypersphere, Mesh, + NoncompactSymmetricSpace, ProductDiscreteSpectrumSpace, SpecialOrthogonal, SpecialUnitary, + SymmetricPositiveDefiniteMatrices, ) from .data import TEST_GRAPH_ADJACENCY, TEST_MESH_PATH @@ -61,6 +64,19 @@ def discrete_spectrum_spaces() -> List[DiscreteSpectrumSpace]: ) +def noncompact_symmetric_spaces() -> List[NoncompactSymmetricSpace]: + return [ + Hyperbolic(2), + Hyperbolic(3), + Hyperbolic(8), + Hyperbolic(9), + SymmetricPositiveDefiniteMatrices(2), + SymmetricPositiveDefiniteMatrices(3), + SymmetricPositiveDefiniteMatrices(6), + SymmetricPositiveDefiniteMatrices(7), + ] + + def np_to_backend(value: B.NPNumeric, backend: str): """ Converts a numpy array to the desired backend. @@ -98,6 +114,11 @@ def np_to_backend(value: B.NPNumeric, backend: str): 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. diff --git a/tests/kernels/test_feature_map_kernel.py b/tests/kernels/test_feature_map_kernel.py new file mode 100644 index 00000000..c3d6d671 --- /dev/null +++ b/tests/kernels/test_feature_map_kernel.py @@ -0,0 +1,135 @@ +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), + - 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() + 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]), + lambda nu, lengthscale, X, X2: kernel.K( + {"nu": nu, "lengthscale": lengthscale}, X, X2 + ), + params["nu"], + params["lengthscale"], + 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 nu, lengthscale, X: kernel.K({"nu": nu, "lengthscale": lengthscale}, X) + - kernel.K({"nu": nu, "lengthscale": lengthscale}, X, X), + params["nu"], + params["lengthscale"], + 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 nu, lengthscale, X: kernel.K_diag( + {"nu": nu, "lengthscale": lengthscale}, X + ) + - B.diag(kernel.K({"nu": nu, "lengthscale": lengthscale}, X)), + params["nu"], + params["lengthscale"], + 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],)), + lambda nu, lengthscale, X: kernel.K_diag( + {"nu": nu, "lengthscale": lengthscale}, X + ), + params["nu"], + params["lengthscale"], + X, + ) From 61c81c48e9b8721881f6be043e9fae897c7d53cd Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Mon, 11 Nov 2024 18:40:21 +0100 Subject: [PATCH 23/35] Remove tests/kernels/test_normalization.py: normalization tests are now part of test_feature_map_kernel and test_matern_karhunenloeve_kernel --- tests/kernels/test_normalization.py | 88 ----------------------------- 1 file changed, 88 deletions(-) delete mode 100644 tests/kernels/test_normalization.py diff --git a/tests/kernels/test_normalization.py b/tests/kernels/test_normalization.py deleted file mode 100644 index 357748f4..00000000 --- a/tests/kernels/test_normalization.py +++ /dev/null @@ -1,88 +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) - num_eigenfns = min(num_eigenfns, A.shape[0]) - - 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) From 0b0a406ab74b2bfd883bd6d88a4741a2debcd20a Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Tue, 12 Nov 2024 13:24:00 +0100 Subject: [PATCH 24/35] Make check_function_with_backend a bit more informative. Add spaces() function that returns the list of all spaces to tests/helper.py. Make tests/spaces/test_basics.py more thorough. --- tests/helper.py | 9 ++++++++- tests/spaces/test_basics.py | 38 ++++++++++++++++++++++++++++++------- 2 files changed, 39 insertions(+), 8 deletions(-) diff --git a/tests/helper.py b/tests/helper.py index 30ae4818..934c3a61 100644 --- a/tests/helper.py +++ b/tests/helper.py @@ -16,6 +16,7 @@ Mesh, NoncompactSymmetricSpace, ProductDiscreteSpectrumSpace, + Space, SpecialOrthogonal, SpecialUnitary, SymmetricPositiveDefiniteMatrices, @@ -77,6 +78,10 @@ def noncompact_symmetric_spaces() -> List[NoncompactSymmetricSpace]: ] +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. @@ -188,7 +193,9 @@ def check_function_with_backend( else: args_casted.append(arg) 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: # we convert `f_output` to numpy array to compare with `result`` if is_bearable(f_output, SparseArray): diff --git a/tests/spaces/test_basics.py b/tests/spaces/test_basics.py index 981447e9..a0e300ee 100644 --- a/tests/spaces/test_basics.py +++ b/tests/spaces/test_basics.py @@ -1,12 +1,36 @@ import inspect +import pytest + import geometric_kernels.spaces -from ..helper import discrete_spectrum_spaces +from ..helper import ( + compact_matrix_lie_groups, + discrete_spectrum_spaces, + noncompact_symmetric_spaces, + product_discrete_spectrum_spaces, + spaces, +) -def test_all_discrete_spectrum_spaces_covered(): - spaces = discrete_spectrum_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 = [ @@ -15,9 +39,7 @@ def test_all_discrete_spectrum_spaces_covered(): if inspect.isclass(cls_obj) ] for cls_name, cls_obj in classes: - if issubclass( - cls_obj, geometric_kernels.spaces.DiscreteSpectrumSpace - ) and not inspect.isabstract(cls_obj): + if issubclass(cls_obj, cls) and not inspect.isabstract(cls_obj): for space in spaces: if isinstance(space, cls_obj): break @@ -25,4 +47,6 @@ def test_all_discrete_spectrum_spaces_covered(): # 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"Space {cls_name} not covered by tests" + assert ( + False + ), f"An instance of the class `{cls_name}` is missing from the list returned by the function `{fun.__name__}`" From 2df9f8e4024e6aee9c342db13da65e0cbcab80f7 Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Tue, 12 Nov 2024 13:27:18 +0100 Subject: [PATCH 25/35] Make random in SpecialOrthogonal and SpecialUnitary compatible with jax and tensorflow by removing item assignments and adding explicit casts --- geometric_kernels/spaces/so.py | 5 +++-- geometric_kernels/spaces/su.py | 13 +++++++++---- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/geometric_kernels/spaces/so.py b/geometric_kernels/spaces/so.py index 847b399a..854a4dc2 100644 --- a/geometric_kernels/spaces/so.py +++ b/geometric_kernels/spaces/so.py @@ -349,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): diff --git a/geometric_kernels/spaces/su.py b/geometric_kernels/spaces/su.py index f529765c..500b74c1 100644 --- a/geometric_kernels/spaces/su.py +++ b/geometric_kernels/spaces/su.py @@ -251,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): From b73c7ed43f8f50e91f78f02fec3597a2db78ba1e Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Tue, 12 Nov 2024 13:29:11 +0100 Subject: [PATCH 26/35] Revise tests/feature_maps/test_feature_maps.py and add tests/sampling/test_samplers.py --- tests/feature_maps/test_feature_maps.py | 118 +++++++++++++----------- tests/sampling/__init__.py | 0 tests/sampling/test_samplers.py | 33 +++++++ 3 files changed, 97 insertions(+), 54 deletions(-) create mode 100644 tests/sampling/__init__.py create mode 100644 tests/sampling/test_samplers.py diff --git a/tests/feature_maps/test_feature_maps.py b/tests/feature_maps/test_feature_maps.py index 926e751a..12684108 100644 --- a/tests/feature_maps/test_feature_maps.py +++ b/tests/feature_maps/test_feature_maps.py @@ -1,69 +1,79 @@ -from pathlib import Path - +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.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 +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.mark.parametrize( - "space_name", ["circle", "hypersphere", "mesh", "graph", "hyperbolic", "spd"] +@pytest.fixture( + params=spaces(), + ids=str, ) -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} +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: - raise ValueError(f"Unknown space {space}") + 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 - 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 = np.random.RandomState() + key, X = space.random(key, 50) - key, points = space.random(key, num_points) + def diff_kern_mats(params, X): + _, embedding = feature_map(X, params) - _, embedding = feature_map(points, params, **kwargs) + kernel_mat = kernel.K(params, X, X) + kernel_mat_alt = B.matmul(embedding, B.T(embedding)) - kernel_mat = kernel.K(params, points, points) - kernel_mat_alt = np.matmul(embedding, embedding.T) + return kernel_mat - kernel_mat_alt - np.testing.assert_allclose(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])), + lambda nu, lengthscale, X: diff_kern_mats( + {"nu": nu, "lengthscale": lengthscale}, X + ), + params["nu"], + params["lengthscale"], + X, + atol=0.1, + ) diff --git a/tests/sampling/__init__.py b/tests/sampling/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/sampling/test_samplers.py b/tests/sampling/test_samplers.py new file mode 100644 index 00000000..90040160 --- /dev/null +++ b/tests/sampling/test_samplers.py @@ -0,0 +1,33 @@ +import lab as B +import numpy as np +import pytest + +from geometric_kernels.sampling import sampler + +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() + key, X = space.random(key, 50) + + # 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), + lambda nu, lengthscale, X: sample_paths( + X, {"nu": nu, "lengthscale": lengthscale}, key=create_random_state(backend) + )[1], + params["nu"], + params["lengthscale"], + X, + compare_to_result=lambda res, f_out: B.shape(f_out) == res, + ) From 57d389b6b165302f99fffbb8b32f9ee9e45e024d Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Tue, 12 Nov 2024 13:36:01 +0100 Subject: [PATCH 27/35] Remove tests/test_dtypes.py: same things are now more thoroughly checked in different places throughout the tests --- tests/test_dtypes.py | 213 ------------------------------------------- 1 file changed, 213 deletions(-) delete mode 100644 tests/test_dtypes.py 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) From 7876a0cf8e2a10fc9a2f53685ae6a2354b42c99b Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Sun, 17 Nov 2024 12:02:43 +0100 Subject: [PATCH 28/35] Move scripts/compute_characters.py from geometric_kernels/utils to scripts/ --- {geometric_kernels/utils => scripts}/compute_characters.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) rename {geometric_kernels/utils => scripts}/compute_characters.py (99%) 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 From bf294c26ef7c4979004fbac1ed9da6223bf0c552 Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Sun, 17 Nov 2024 20:00:57 +0100 Subject: [PATCH 29/35] Continuing to revise tests. Many changes, see the list below. - Implemented some alternative formulas for the heat kernel on `Hyperbolic` and `SymmetricPositiveDefiniteMatrices` for testing, along with the respective tests. - Revised and refactored the tests for `ProductDiscreteSpectrumSpace` and `ProductGeometricKernel`. - Changed `params_to_params_list` function to take the number of factors as an explicit parameter. `ProductGeometricKernel` can now take 1-d length scale and nu parameters. - Rewritten `random` for `Hyperbolic` and `SymmetricPositiveDefiniteMatrices` to actually use the provided random key. - Moved the implementation of `Hyperbolic.distance` and `Hyperbolic.inner_product` to utils.manifold_utils (the functions `hyperbolic_distance` and `minkowski_inner_product`). - Moved MaternKarhunenLoeveKernel import inside __init__ in feature_maps/random_phase.py to avoid circular imports. - Added an informative error message in case `check_function_with_backend` fails because of shape mismatch. - Added a workaround for a bug in lab (https://github.com/wesselb/lab/issues/21) - it was only an issue for one of the tests and couldn't have affected the main code. - Moved all kernel formulas from `utils.special_functions` to a new subpackage `utils.kernel_formulas`. - Updated bib entries for the "Stationary Kernels and Gaussian Processes on Lie Groups and their Homogeneous Spaces" papers which have been recently accepted by JMLR. --- docs/references.bib | 33 ++- docs/theory/addition_theorem.rst | 2 +- docs/theory/symmetric.rst | 4 +- .../feature_maps/probability_densities.py | 4 +- .../feature_maps/random_phase.py | 3 +- geometric_kernels/kernels/product.py | 4 +- geometric_kernels/spaces/base.py | 2 +- geometric_kernels/spaces/hyperbolic.py | 95 ++------- geometric_kernels/spaces/lie_groups.py | 6 +- geometric_kernels/spaces/so.py | 4 +- geometric_kernels/spaces/spd.py | 15 +- geometric_kernels/spaces/su.py | 4 +- .../utils/kernel_formulas/__init__.py | 14 ++ .../utils/kernel_formulas/euclidean.py | 103 +++++++++ .../utils/kernel_formulas/hyperbolic.py | 199 ++++++++++++++++++ .../utils/kernel_formulas/hypercube_graph.py | 52 +++++ .../utils/kernel_formulas/spd.py | 127 +++++++++++ geometric_kernels/utils/manifold_utils.py | 79 +++++++ geometric_kernels/utils/product.py | 15 +- geometric_kernels/utils/special_functions.py | 136 ------------ notebooks/Hyperbolic.ipynb | 13 +- notebooks/SPD.ipynb | 13 +- notebooks/SpecialOrthogonal.ipynb | 13 +- notebooks/SpecialUnitary.ipynb | 13 +- .../other/Hyperbolic Approximations.ipynb | 13 +- notebooks/other/SPD Approximations.ipynb | 13 +- tests/helper.py | 3 + .../test_matern_karhunenloeve_kernel.py | 11 +- tests/kernels/test_product_kernel.py | 51 +++++ tests/sampling/test_samplers.py | 1 + tests/spaces/test_circle.py | 18 +- tests/spaces/test_hyperbolic.py | 77 ++++--- tests/spaces/test_hypercube_graph.py | 2 +- .../test_product_discrete_spectrum_space.py | 61 ++++++ tests/spaces/test_spd.py | 43 ++++ tests/test_product.py | 189 ----------------- tests/utils/test_kernel_formulas.py | 67 ++++++ tests/utils/test_manifold_utils.py | 32 +++ tests/utils/test_special_functions.py | 62 +----- 39 files changed, 1031 insertions(+), 565 deletions(-) create mode 100644 geometric_kernels/utils/kernel_formulas/__init__.py create mode 100644 geometric_kernels/utils/kernel_formulas/euclidean.py create mode 100644 geometric_kernels/utils/kernel_formulas/hyperbolic.py create mode 100644 geometric_kernels/utils/kernel_formulas/hypercube_graph.py create mode 100644 geometric_kernels/utils/kernel_formulas/spd.py create mode 100644 tests/kernels/test_product_kernel.py create mode 100644 tests/spaces/test_product_discrete_spectrum_space.py create mode 100644 tests/spaces/test_spd.py delete mode 100644 tests/test_product.py create mode 100644 tests/utils/test_kernel_formulas.py create mode 100644 tests/utils/test_manifold_utils.py 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/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..d4b214b7 100644 --- a/geometric_kernels/kernels/product.py +++ b/geometric_kernels/kernels/product.py @@ -116,7 +116,7 @@ def K(self, params: Dict[str, B.Numeric], X, X2=None, **kwargs) -> B.Numeric: 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) + params_list = params_to_params_list(len(self.kernels), params) return B.prod( B.stack( @@ -131,7 +131,7 @@ 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) + params_list = params_to_params_list(len(self.kernels), params) return B.prod( B.stack( diff --git a/geometric_kernels/spaces/base.py b/geometric_kernels/spaces/base.py index 6d192bf2..b8ae32be 100644 --- a/geometric_kernels/spaces/base.py +++ b/geometric_kernels/spaces/base.py @@ -170,7 +170,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/hyperbolic.py b/geometric_kernels/spaces/hyperbolic.py index 4e0a336f..f03d63db 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,12 +41,12 @@ 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): @@ -67,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) @@ -208,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. @@ -219,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): """ diff --git a/geometric_kernels/spaces/lie_groups.py b/geometric_kernels/spaces/lie_groups.py index 6e3877e8..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. @@ -196,7 +196,7 @@ 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) @@ -217,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. diff --git a/geometric_kernels/spaces/so.py b/geometric_kernels/spaces/so.py index 854a4dc2..b36c90cf 100644 --- a/geometric_kernels/spaces/so.py +++ b/geometric_kernels/spaces/so.py @@ -193,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. @@ -261,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): diff --git a/geometric_kernels/spaces/spd.py b/geometric_kernels/spaces/spd.py index 3135d98d..142ffe2a 100644 --- a/geometric_kernels/spaces/spd.py +++ b/geometric_kernels/spaces/spd.py @@ -43,12 +43,12 @@ 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): @@ -120,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. @@ -131,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): diff --git a/geometric_kernels/spaces/su.py b/geometric_kernels/spaces/su.py index 500b74c1..004844bb 100644 --- a/geometric_kernels/spaces/su.py +++ b/geometric_kernels/spaces/su.py @@ -108,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. @@ -177,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): 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..779480e1 --- /dev/null +++ b/geometric_kernels/utils/kernel_formulas/hyperbolic.py @@ -0,0 +1,199 @@ +""" +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 + # print(f"norm = {normalization}, dist={rho}, cur_result = {cur_result}, result = {result[i, j]}") + + 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..eb8d506e 100644 --- a/geometric_kernels/utils/product.py +++ b/geometric_kernels/utils/product.py @@ -4,18 +4,27 @@ from beartype.typing import Dict, List -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): diff --git a/geometric_kernels/utils/special_functions.py b/geometric_kernels/utils/special_functions.py index 69add35c..3fa9d6ee 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,135 +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 - - -def 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 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 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 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/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/tests/helper.py b/tests/helper.py index 934c3a61..5c28b940 100644 --- a/tests/helper.py +++ b/tests/helper.py @@ -202,6 +202,9 @@ def check_function_with_backend( 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( diff --git a/tests/kernels/test_matern_karhunenloeve_kernel.py b/tests/kernels/test_matern_karhunenloeve_kernel.py index fcee3cf5..b2b2fcdc 100644 --- a/tests/kernels/test_matern_karhunenloeve_kernel.py +++ b/tests/kernels/test_matern_karhunenloeve_kernel.py @@ -194,10 +194,13 @@ def test_normalize(inputs, backend): # Check that the average variance of the kernel is 1. check_function_with_backend( backend, - np.array(1.0), - lambda nu, lengthscale, X: B.mean( - kernel.K_diag({"nu": nu, "lengthscale": lengthscale}, X), squeeze=False - ), + np.array([1.0]), + lambda nu, lengthscale, X: B.reshape( + B.mean( + kernel.K_diag({"nu": nu, "lengthscale": lengthscale}, X), squeeze=False + ), + 1, + ), # the reshape shields from a bug in lab present at least up to version 1.6.6 params["nu"], params["lengthscale"], X, diff --git a/tests/kernels/test_product_kernel.py b/tests/kernels/test_product_kernel.py new file mode 100644 index 00000000..773d4470 --- /dev/null +++ b/tests/kernels/test_product_kernel.py @@ -0,0 +1,51 @@ +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", "jax"]) +# @pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_heat_kernel_is_product_of_heat_kernels( + factor1, factor2, nu, lengthscale, backend +): + key = np.random.RandomState() + 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} + + # common_dtype = B.promote_dtypes(B.dtype(xs_factor1), B.dtype(xs_factor2)) + # xs_product = make_product([B.cast(common_dtype, xs_factor1), B.cast(common_dtype, xs_factor2)]) + 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/sampling/test_samplers.py b/tests/sampling/test_samplers.py index 90040160..f2493f54 100644 --- a/tests/sampling/test_samplers.py +++ b/tests/sampling/test_samplers.py @@ -4,6 +4,7 @@ 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 diff --git a/tests/spaces/test_circle.py b/tests/spaces/test_circle.py index f6e6689c..dda7e95d 100644 --- a/tests/spaces/test_circle.py +++ b/tests/spaces/test_circle.py @@ -4,11 +4,11 @@ from geometric_kernels.kernels import MaternGeometricKernel from geometric_kernels.spaces.circle import Circle -from geometric_kernels.utils.special_functions import ( - matern_12_kernel, - matern_32_kernel, - matern_52_kernel, - rbf_kernel, +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 @@ -18,13 +18,13 @@ @pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) def test_equivalence_kernel(nu, atol, backend): if nu == 0.5: - analytic_kernel = matern_12_kernel + analytic_kernel = euclidean_matern_12_kernel elif nu == 1.5: - analytic_kernel = matern_32_kernel + analytic_kernel = euclidean_matern_32_kernel elif nu == 2.5: - analytic_kernel = matern_52_kernel + analytic_kernel = euclidean_matern_52_kernel elif nu == np.inf: - analytic_kernel = rbf_kernel + 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)) diff --git a/tests/spaces/test_hyperbolic.py b/tests/spaces/test_hyperbolic.py index 79e2d0ee..8548ed3c 100644 --- a/tests/spaces/test_hyperbolic.py +++ b/tests/spaces/test_hyperbolic.py @@ -1,28 +1,51 @@ +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", [0.1, 1.0, 10.0]) +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_equivalence_kernel_odd(dim, lengthscale, backend): + space = Hyperbolic(dim) + + key = np.random.RandomState() + 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)) + + # 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, + lambda nu, lengthscale, X, X2: kernel.K( + {"nu": nu, "lengthscale": lengthscale}, X, X2 + ), + np.array([np.inf]), + np.array([lengthscale]), + X, + X2, + compare_to_result=lambda res, f_out: np.linalg.norm(res - B.to_numpy(f_out)) + / np.sqrt(res.shape[0] * res.shape[1]) + < 1e-1, + ) diff --git a/tests/spaces/test_hypercube_graph.py b/tests/spaces/test_hypercube_graph.py index 73296894..745c654f 100644 --- a/tests/spaces/test_hypercube_graph.py +++ b/tests/spaces/test_hypercube_graph.py @@ -5,7 +5,7 @@ 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.kernel_formulas import hypercube_graph_heat_kernel from ..helper import check_function_with_backend 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..9c984e27 --- /dev/null +++ b/tests/spaces/test_product_discrete_spectrum_space.py @@ -0,0 +1,61 @@ +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 + +# _TRUNC_LEVEL = 128 +# _GRID_SIZE = 3 +_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", "jax"]) +# @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() + 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} + + # common_dtype = B.promote_dtypes(B.dtype(xs_factor1), B.dtype(xs_factor2)) + # xs_product = make_product([B.cast(common_dtype, xs_factor1), B.cast(common_dtype, xs_factor2)]) + 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..02cde3fd --- /dev/null +++ b/tests/spaces/test_spd.py @@ -0,0 +1,43 @@ +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", [0.5, 1.0, 10.0]) +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_equivalence_kernel(lengthscale, backend): + space = SymmetricPositiveDefiniteMatrices(2) + + key = np.random.RandomState() + 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)) + + # 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, + lambda nu, lengthscale, X, X2: kernel.K( + {"nu": nu, "lengthscale": lengthscale}, X, X2 + ), + np.array([np.inf]), + np.array([lengthscale]), + X, + X2, + compare_to_result=lambda res, f_out: np.linalg.norm(res - B.to_numpy(f_out)) + / np.sqrt(res.shape[0] * res.shape[1]) + < 1e-1, + ) diff --git a/tests/test_product.py b/tests/test_product.py deleted file mode 100644 index bb345c37..00000000 --- a/tests/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/utils/test_kernel_formulas.py b/tests/utils/test_kernel_formulas.py new file mode 100644 index 00000000..8013f687 --- /dev/null +++ b/tests/utils/test_kernel_formulas.py @@ -0,0 +1,67 @@ +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() + 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, :], + ) diff --git a/tests/utils/test_manifold_utils.py b/tests/utils/test_manifold_utils.py new file mode 100644 index 00000000..56604570 --- /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() + 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..d8e57a23 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, ) @@ -110,60 +107,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, :], - ) From a46ff93abee65f90733a6eb6f2d89164ab2989f1 Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Sun, 17 Nov 2024 21:57:09 +0100 Subject: [PATCH 30/35] Change parameters in tests/spaces/test_hyperbolic.py and tests/spaces/test_spd.py to speed up tests and make them more robust in the light of numerical instabilities --- tests/spaces/test_hyperbolic.py | 10 +++++++--- tests/spaces/test_spd.py | 2 +- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/tests/spaces/test_hyperbolic.py b/tests/spaces/test_hyperbolic.py index 8548ed3c..82030e2e 100644 --- a/tests/spaces/test_hyperbolic.py +++ b/tests/spaces/test_hyperbolic.py @@ -11,11 +11,15 @@ from ..helper import check_function_with_backend, create_random_state +# from line_profiler import profile + +# @profile @pytest.mark.parametrize("dim", [2, 3, 5, 7]) -@pytest.mark.parametrize("lengthscale", [0.1, 1.0, 10.0]) -@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) -def test_equivalence_kernel_odd(dim, lengthscale, backend): +@pytest.mark.parametrize("lengthscale", [2.0]) +@pytest.mark.parametrize("backend", ["tensorflow"]) +# @pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +def test_equivalence_kernel(dim, lengthscale, backend): space = Hyperbolic(dim) key = np.random.RandomState() diff --git a/tests/spaces/test_spd.py b/tests/spaces/test_spd.py index 02cde3fd..bf1327a3 100644 --- a/tests/spaces/test_spd.py +++ b/tests/spaces/test_spd.py @@ -9,7 +9,7 @@ from ..helper import check_function_with_backend, create_random_state -@pytest.mark.parametrize("lengthscale", [0.5, 1.0, 10.0]) +@pytest.mark.parametrize("lengthscale", [2.0]) @pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) def test_equivalence_kernel(lengthscale, backend): space = SymmetricPositiveDefiniteMatrices(2) From c050d5345904e515497e00d779e23938c7f7de6c Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Mon, 18 Nov 2024 12:17:32 +0100 Subject: [PATCH 31/35] Fix typing issues in `ProductGeometricKernel` and `ProductDiscreteSpectrumSpace` that were arising under torch and tensorflow when you mix complex and real data types. More details below. --- geometric_kernels/kernels/product.py | 13 ++++- geometric_kernels/lab_extras/extras.py | 34 +++++++++++ geometric_kernels/lab_extras/jax/extras.py | 36 +++++++++--- geometric_kernels/lab_extras/numpy/extras.py | 24 ++++++-- .../lab_extras/tensorflow/extras.py | 22 +++++-- geometric_kernels/lab_extras/torch/extras.py | 22 +++++-- geometric_kernels/spaces/base.py | 18 +++++- geometric_kernels/spaces/circle.py | 8 +++ geometric_kernels/spaces/graph.py | 8 +++ geometric_kernels/spaces/hyperbolic.py | 8 +++ geometric_kernels/spaces/hypercube_graph.py | 8 +++ geometric_kernels/spaces/hypersphere.py | 8 +++ geometric_kernels/spaces/mesh.py | 8 +++ geometric_kernels/spaces/product.py | 57 +++++++++++++------ geometric_kernels/spaces/so.py | 8 +++ geometric_kernels/spaces/spd.py | 8 +++ geometric_kernels/spaces/su.py | 8 +++ geometric_kernels/utils/product.py | 21 +++++-- tests/kernels/test_product_kernel.py | 7 +-- tests/spaces/test_hyperbolic.py | 6 +- .../test_product_discrete_spectrum_space.py | 7 +-- 21 files changed, 275 insertions(+), 64 deletions(-) diff --git a/geometric_kernels/kernels/product.py b/geometric_kernels/kernels/product.py index d4b214b7..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,8 +115,12 @@ 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) + 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( @@ -130,7 +135,9 @@ 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) + 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( 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..efe61bdc 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,11 @@ 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) + # if B.dtype(reference) != jnp.float64: + # return jnp.complex64 + # else: + # return jnp.complex128 @dispatch @@ -244,4 +249,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..adc96f5f 100644 --- a/geometric_kernels/lab_extras/numpy/extras.py +++ b/geometric_kernels/lab_extras/numpy/extras.py @@ -144,10 +144,11 @@ 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) + # if reference.dtype != np.float64: + # return np.complex64 + # else: + # return np.complex128 @dispatch @@ -239,4 +240,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/tensorflow/extras.py b/geometric_kernels/lab_extras/tensorflow/extras.py index 066732e2..aea2414a 100644 --- a/geometric_kernels/lab_extras/tensorflow/extras.py +++ b/geometric_kernels/lab_extras/tensorflow/extras.py @@ -170,10 +170,11 @@ 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) + # if B.dtype(reference) != tf.float64: + # return tf.complex64 + # else: + # return tf.complex128 @dispatch @@ -257,3 +258,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 823e563a..20f8feae 100644 --- a/geometric_kernels/lab_extras/torch/extras.py +++ b/geometric_kernels/lab_extras/torch/extras.py @@ -171,10 +171,11 @@ 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) + # if B.dtype(reference) != torch.double: + # return torch.cfloat + # else: + # return torch.cdouble @dispatch @@ -258,3 +259,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 b8ae32be..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 diff --git a/geometric_kernels/spaces/circle.py b/geometric_kernels/spaces/circle.py index b88f448e..d80d4d51 100644 --- a/geometric_kernels/spaces/circle.py +++ b/geometric_kernels/spaces/circle.py @@ -188,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/graph.py b/geometric_kernels/spaces/graph.py index 94cf71da..bfdf37f5 100644 --- a/geometric_kernels/spaces/graph.py +++ b/geometric_kernels/spaces/graph.py @@ -190,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 f03d63db..e0417f92 100644 --- a/geometric_kernels/spaces/hyperbolic.py +++ b/geometric_kernels/spaces/hyperbolic.py @@ -173,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 1c315523..2034fead 100644 --- a/geometric_kernels/spaces/hypercube_graph.py +++ b/geometric_kernels/spaces/hypercube_graph.py @@ -272,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 39c22b85..5aa1870a 100644 --- a/geometric_kernels/spaces/hypersphere.py +++ b/geometric_kernels/spaces/hypersphere.py @@ -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/mesh.py b/geometric_kernels/spaces/mesh.py index a1cb37e5..122806f7 100644 --- a/geometric_kernels/spaces/mesh.py +++ b/geometric_kernels/spaces/mesh.py @@ -217,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 74d2b3ee..d475d07a 100644 --- a/geometric_kernels/spaces/product.py +++ b/geometric_kernels/spaces/product.py @@ -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 + ) factor_eigenfunction_values = [ eigenfunction(X, **kwargs) # [N, Js], Js different for each factor @@ -281,21 +289,27 @@ 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 + ) + Xs2 = project_product( + X2, self.dimension_indices, self.element_shapes, self.element_dtypes + ) + + factor_phi_products = [ + take_along_axis( + eigenfunction.phi_product(X1, X2, **kwargs), + from_numpy(X1, self.eigenindicies[None, None, :, s]), + -1, + ) + 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( - *[ - take_along_axis( - eigenfunction.phi_product(X1, X2, **kwargs), - from_numpy(X1, self.eigenindicies[None, None, :, s]), - -1, - ) - for s, (eigenfunction, X1, X2) in enumerate( - zip(self.eigenfunctions, Xs, Xs2) - ) - ], - axis=-1, + *[B.cast(common_dtype, x) for x in factor_phi_products], axis=-1 ) # [N, N2, L, S] prod_phis = B.prod(phis, axis=-1) # [N, N2, L, S] -> [N, N2, L] @@ -303,7 +317,9 @@ def phi_product( 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( *[ @@ -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 b36c90cf..6ea1969a 100644 --- a/geometric_kernels/spaces/so.py +++ b/geometric_kernels/spaces/so.py @@ -360,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 142ffe2a..13ec3952 100644 --- a/geometric_kernels/spaces/spd.py +++ b/geometric_kernels/spaces/spd.py @@ -147,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 004844bb..5fc405df 100644 --- a/geometric_kernels/spaces/su.py +++ b/geometric_kernels/spaces/su.py @@ -270,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/product.py b/geometric_kernels/utils/product.py index eb8d506e..58baded3 100644 --- a/geometric_kernels/utils/product.py +++ b/geometric_kernels/utils/product.py @@ -3,6 +3,8 @@ import lab as B from beartype.typing import Dict, List +from geometric_kernels.lab_extras import smart_cast + def params_to_params_list( number_of_factors: int, params: Dict[str, B.Numeric] @@ -52,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. @@ -72,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 @@ -80,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/tests/kernels/test_product_kernel.py b/tests/kernels/test_product_kernel.py index 773d4470..7fb39ef9 100644 --- a/tests/kernels/test_product_kernel.py +++ b/tests/kernels/test_product_kernel.py @@ -12,11 +12,8 @@ "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", "jax"]) -# @pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) -def test_heat_kernel_is_product_of_heat_kernels( - factor1, factor2, nu, lengthscale, backend -): +@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() key, xs_factor1 = factor1.random(key, 10) key, xs_factor2 = factor2.random(key, 10) diff --git a/tests/spaces/test_hyperbolic.py b/tests/spaces/test_hyperbolic.py index 82030e2e..c76a0d18 100644 --- a/tests/spaces/test_hyperbolic.py +++ b/tests/spaces/test_hyperbolic.py @@ -11,14 +11,10 @@ from ..helper import check_function_with_backend, create_random_state -# from line_profiler import profile - -# @profile @pytest.mark.parametrize("dim", [2, 3, 5, 7]) @pytest.mark.parametrize("lengthscale", [2.0]) -@pytest.mark.parametrize("backend", ["tensorflow"]) -# @pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) def test_equivalence_kernel(dim, lengthscale, backend): space = Hyperbolic(dim) diff --git a/tests/spaces/test_product_discrete_spectrum_space.py b/tests/spaces/test_product_discrete_spectrum_space.py index 9c984e27..2aa2910b 100644 --- a/tests/spaces/test_product_discrete_spectrum_space.py +++ b/tests/spaces/test_product_discrete_spectrum_space.py @@ -11,8 +11,6 @@ from ..helper import check_function_with_backend -# _TRUNC_LEVEL = 128 -# _GRID_SIZE = 3 _NUM_LEVELS = 20 @@ -20,8 +18,7 @@ "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", "jax"]) -# @pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) +@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 @@ -38,8 +35,6 @@ def test_heat_kernel_is_product_of_heat_kernels(factor1, factor2, lengthscale, b def K_diff(nu, lengthscale, xs_factor1, xs_factor2): params = {"nu": nu, "lengthscale": lengthscale} - # common_dtype = B.promote_dtypes(B.dtype(xs_factor1), B.dtype(xs_factor2)) - # xs_product = make_product([B.cast(common_dtype, xs_factor1), B.cast(common_dtype, xs_factor2)]) xs_product = make_product([xs_factor1, xs_factor2]) K_product = kernel_product.K(params, xs_product, xs_product) From d32bb350ea6229cb60f4b9c05b7b9eb92dbd3fc2 Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Mon, 18 Nov 2024 13:45:38 +0100 Subject: [PATCH 32/35] Add pytest-cov to test_requirements.txt and flags --cov --cov-report=xml in make test --- Makefile | 2 +- test_requirements.txt | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) 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/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 From edb6169da1cb9cb84d5197acc8028f42cf1e85c9 Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Thu, 21 Nov 2024 22:44:22 +0100 Subject: [PATCH 33/35] Incorporate the feedback from stoprightthere's review --- geometric_kernels/lab_extras/jax/extras.py | 4 -- geometric_kernels/lab_extras/numpy/extras.py | 4 -- .../lab_extras/tensorflow/extras.py | 4 -- geometric_kernels/lab_extras/torch/extras.py | 4 -- geometric_kernels/spaces/product.py | 22 +++---- .../utils/kernel_formulas/hyperbolic.py | 1 - tests/feature_maps/test_feature_maps.py | 2 +- tests/feature_maps/test_student_t_sample.py | 2 +- tests/kernels/test_feature_map_kernel.py | 36 ++++++----- .../test_matern_karhunenloeve_kernel.py | 59 +++++++++++-------- tests/kernels/test_product_kernel.py | 4 +- tests/sampling/test_samplers.py | 11 ++-- tests/spaces/test_circle.py | 7 ++- tests/spaces/test_eigenfunctions.py | 4 +- tests/spaces/test_hyperbolic.py | 9 +-- tests/spaces/test_hypercube_graph.py | 9 +-- tests/spaces/test_lie_groups.py | 4 +- .../test_product_discrete_spectrum_space.py | 2 +- tests/spaces/test_spd.py | 9 +-- tests/utils/test_kernel_formulas.py | 15 ++--- tests/utils/test_manifold_utils.py | 2 +- tests/utils/test_special_functions.py | 13 ++-- 22 files changed, 119 insertions(+), 108 deletions(-) diff --git a/geometric_kernels/lab_extras/jax/extras.py b/geometric_kernels/lab_extras/jax/extras.py index efe61bdc..5743ef3e 100644 --- a/geometric_kernels/lab_extras/jax/extras.py +++ b/geometric_kernels/lab_extras/jax/extras.py @@ -160,10 +160,6 @@ def complex_like(reference: B.JAXNumeric): Return `complex` dtype of a backend based on the reference. """ return B.promote_dtypes(jnp.complex64, reference.dtype) - # if B.dtype(reference) != jnp.float64: - # return jnp.complex64 - # else: - # return jnp.complex128 @dispatch diff --git a/geometric_kernels/lab_extras/numpy/extras.py b/geometric_kernels/lab_extras/numpy/extras.py index adc96f5f..01c69168 100644 --- a/geometric_kernels/lab_extras/numpy/extras.py +++ b/geometric_kernels/lab_extras/numpy/extras.py @@ -145,10 +145,6 @@ def complex_like(reference: B.NPNumeric): Return `complex` dtype of a backend based on the reference. """ return B.promote_dtypes(np.complex64, reference.dtype) - # if reference.dtype != np.float64: - # return np.complex64 - # else: - # return np.complex128 @dispatch diff --git a/geometric_kernels/lab_extras/tensorflow/extras.py b/geometric_kernels/lab_extras/tensorflow/extras.py index aea2414a..594be7f9 100644 --- a/geometric_kernels/lab_extras/tensorflow/extras.py +++ b/geometric_kernels/lab_extras/tensorflow/extras.py @@ -171,10 +171,6 @@ def complex_like(reference: B.TFNumeric): Return `complex` dtype of a backend based on the reference. """ return B.promote_dtypes(tf.complex64, reference.dtype) - # if B.dtype(reference) != tf.float64: - # return tf.complex64 - # else: - # return tf.complex128 @dispatch diff --git a/geometric_kernels/lab_extras/torch/extras.py b/geometric_kernels/lab_extras/torch/extras.py index 20f8feae..b16a5759 100644 --- a/geometric_kernels/lab_extras/torch/extras.py +++ b/geometric_kernels/lab_extras/torch/extras.py @@ -172,10 +172,6 @@ def complex_like(reference: B.TorchNumeric): Return `complex` dtype of a backend based on the reference. """ return B.promote_dtypes(torch.cfloat, reference.dtype) - # if B.dtype(reference) != torch.double: - # return torch.cfloat - # else: - # return torch.cdouble @dispatch diff --git a/geometric_kernels/spaces/product.py b/geometric_kernels/spaces/product.py index d475d07a..c2b7de1c 100644 --- a/geometric_kernels/spaces/product.py +++ b/geometric_kernels/spaces/product.py @@ -248,7 +248,7 @@ def __call__(self, X: B.Numeric, **kwargs) -> B.Numeric: """ 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 @@ -264,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] @@ -291,17 +291,17 @@ def phi_product( X2 = X 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), + 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) ) @@ -324,16 +324,16 @@ def phi_product_diag(self, X: B.Numeric, **kwargs): phis = B.stack( *[ take_along_axis( - eigenfunction.phi_product_diag(X1, **kwargs), + 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, L, S] - prod_phis = B.prod(phis, axis=-1) # [N, L, S] -> [N, L] + prod_phis = B.prod(phis, axis=-1) # [N, L] return prod_phis diff --git a/geometric_kernels/utils/kernel_formulas/hyperbolic.py b/geometric_kernels/utils/kernel_formulas/hyperbolic.py index 779480e1..575c0612 100644 --- a/geometric_kernels/utils/kernel_formulas/hyperbolic.py +++ b/geometric_kernels/utils/kernel_formulas/hyperbolic.py @@ -194,6 +194,5 @@ def hyperbolic_heat_kernel_even( rho = hyperbolic_distance(x, x2).squeeze() cur_result = _hyperbolic_heat_kernel_2d_unnormalized(t, rho) result[i, j] = cur_result / normalization - # print(f"norm = {normalization}, dist={rho}, cur_result = {cur_result}, result = {result[i, j]}") return result diff --git a/tests/feature_maps/test_feature_maps.py b/tests/feature_maps/test_feature_maps.py index 12684108..1bf0d235 100644 --- a/tests/feature_maps/test_feature_maps.py +++ b/tests/feature_maps/test_feature_maps.py @@ -53,7 +53,7 @@ def test_feature_map_approximates_kernel(backend, feature_map_and_friends): params = kernel.init_params() - key = np.random.RandomState() + key = np.random.RandomState(0) key, X = space.random(key, 50) def diff_kern_mats(params, X): diff --git a/tests/feature_maps/test_student_t_sample.py b/tests/feature_maps/test_student_t_sample.py index 5b262f15..9a1051d7 100644 --- a/tests/feature_maps/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/kernels/test_feature_map_kernel.py b/tests/kernels/test_feature_map_kernel.py index c3d6d671..08fb8f81 100644 --- a/tests/kernels/test_feature_map_kernel.py +++ b/tests/kernels/test_feature_map_kernel.py @@ -21,7 +21,7 @@ def inputs(request): """ Returns a tuple (space, num_features, feature_map, X, X2) where: - space = request.param, - - num_features = default_num(space), + - 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, @@ -30,7 +30,7 @@ def inputs(request): num_features = min(default_num(space), 15) feature_map = default_feature_map(space=space, num=num_features) - key = np.random.RandomState() + 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) @@ -63,13 +63,14 @@ def test_K(inputs, backend, normalize, kernel): _, _, _, X, X2 = inputs params = kernel.init_params() + def kern(nu, lengthscale, X, X2): + return kernel.K({"nu": nu, "lengthscale": lengthscale}, X, X2) + # 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]), - lambda nu, lengthscale, X, X2: kernel.K( - {"nu": nu, "lengthscale": lengthscale}, X, X2 - ), + kern, params["nu"], params["lengthscale"], X, @@ -84,12 +85,16 @@ def test_K_one_param(inputs, backend, normalize, kernel): _, _, _, X, _ = inputs params = kernel.init_params() + def diff(nu, lengthscale, X): + return kernel.K({"nu": nu, "lengthscale": lengthscale}, X) - kernel.K( + {"nu": nu, "lengthscale": lengthscale}, X, X + ) + # 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 nu, lengthscale, X: kernel.K({"nu": nu, "lengthscale": lengthscale}, X) - - kernel.K({"nu": nu, "lengthscale": lengthscale}, X, X), + diff, params["nu"], params["lengthscale"], X, @@ -102,14 +107,16 @@ def test_K_diag(inputs, backend, normalize, kernel): _, _, _, X, _ = inputs params = kernel.init_params() + def diff(nu, lengthscale, X): + return kernel.K_diag({"nu": nu, "lengthscale": lengthscale}, X) - B.diag( + kernel.K({"nu": nu, "lengthscale": lengthscale}, X) + ) + # Check that kernel.K_diag coincides with the diagonal of kernel.K. check_function_with_backend( backend, np.zeros((X.shape[0],)), - lambda nu, lengthscale, X: kernel.K_diag( - {"nu": nu, "lengthscale": lengthscale}, X - ) - - B.diag(kernel.K({"nu": nu, "lengthscale": lengthscale}, X)), + diff, params["nu"], params["lengthscale"], X, @@ -122,13 +129,14 @@ def test_normalize(inputs, backend, kernel): params = kernel.init_params() + def kern_diag(nu, lengthscale, X): + return kernel.K_diag({"nu": nu, "lengthscale": lengthscale}, X) + # Check that the variance of the kernel is constant 1. check_function_with_backend( backend, np.ones((X.shape[0],)), - lambda nu, lengthscale, X: kernel.K_diag( - {"nu": nu, "lengthscale": lengthscale}, X - ), + kern_diag, params["nu"], params["lengthscale"], X, diff --git a/tests/kernels/test_matern_karhunenloeve_kernel.py b/tests/kernels/test_matern_karhunenloeve_kernel.py index b2b2fcdc..62a23e8a 100644 --- a/tests/kernels/test_matern_karhunenloeve_kernel.py +++ b/tests/kernels/test_matern_karhunenloeve_kernel.py @@ -31,7 +31,7 @@ def inputs(request): num_levels = default_num(space) kernel = MaternKarhunenLoeveKernel(space, num_levels, normalize=normalize) - key = np.random.RandomState() + 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) @@ -61,13 +61,14 @@ def test_eigenvalues_shape(inputs, backend): _, num_levels, kernel, _, _ = inputs params = kernel.init_params() + def eigenvalues(nu, lengthscale): + return kernel.eigenvalues({"nu": nu, "lengthscale": lengthscale}) + # Check that the eigenvalues have appropriate shape. check_function_with_backend( backend, (num_levels, 1), - lambda nu, lengthscale: kernel.eigenvalues( - {"nu": nu, "lengthscale": lengthscale} - ), + eigenvalues, params["nu"], params["lengthscale"], compare_to_result=lambda res, f_out: f_out.shape == res, @@ -79,13 +80,14 @@ def test_eigenvalues_positive(inputs, backend): _, _, kernel, _, _ = inputs params = kernel.init_params() + def eigenvalues(nu, lengthscale): + return kernel.eigenvalues({"nu": nu, "lengthscale": lengthscale}) + # Check that the eigenvalues are nonnegative. check_function_with_backend( backend, None, - lambda nu, lengthscale: kernel.eigenvalues( - {"nu": nu, "lengthscale": lengthscale} - ), + eigenvalues, params["nu"], params["lengthscale"], compare_to_result=lambda _, f_out: np.all(B.to_numpy(f_out) >= 0), @@ -97,13 +99,14 @@ def test_eigenvalues_ordered(inputs, backend): _, _, kernel, _, _ = inputs params = kernel.init_params() + def eigenvalues(nu, lengthscale): + return kernel.eigenvalues({"nu": nu, "lengthscale": lengthscale}) + # Check that the eigenvalues are sorted in descending order. check_function_with_backend( backend, None, - lambda nu, lengthscale: kernel.eigenvalues( - {"nu": nu, "lengthscale": lengthscale} - ), + eigenvalues, params["nu"], params["lengthscale"], compare_to_result=lambda _, f_out: np.all( @@ -121,14 +124,15 @@ def test_K(inputs, backend): assert result.shape == (X.shape[0], X2.shape[0]), "K has incorrect shape" + def kern(nu, lengthscale, X, X2): + return kernel.K({"nu": nu, "lengthscale": lengthscale}, X, X2) + if backend != "numpy": # Check that kernel.K computed using `backend` coincides with the numpy result. check_function_with_backend( backend, result, - lambda nu, lengthscale, X, X2: kernel.K( - {"nu": nu, "lengthscale": lengthscale}, X, X2 - ), + kern, params["nu"], params["lengthscale"], X, @@ -143,11 +147,14 @@ def test_K_one_param(inputs, backend): result = kernel.K(params, X, X) + def kern(nu, lengthscale, X): + return kernel.K({"nu": nu, "lengthscale": lengthscale}, X) + # Check that kernel.K(X) coincides with kernel.K(X, X). check_function_with_backend( backend, result, - lambda nu, lengthscale, X: kernel.K({"nu": nu, "lengthscale": lengthscale}, X), + kern, params["nu"], params["lengthscale"], X, @@ -164,13 +171,14 @@ def test_K_diag(inputs, backend): assert result.shape == (X.shape[0],), "The diagonal has incorrect shape" + def kern_diag(nu, lengthscale, X): + return kernel.K_diag({"nu": nu, "lengthscale": lengthscale}, X) + # Check that kernel.K_diag coincides with the diagonal of kernel.K. check_function_with_backend( backend, result, - lambda nu, lengthscale, X: kernel.K_diag( - {"nu": nu, "lengthscale": lengthscale}, X - ), + kern_diag, params["nu"], params["lengthscale"], X, @@ -186,21 +194,24 @@ def test_normalize(inputs, backend): pytest.skip("No need to check normalization for an unnormalized kernel") params = kernel.init_params() - key = np.random.RandomState() + 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 - # Check that the average variance of the kernel is 1. - check_function_with_backend( - backend, - np.array([1.0]), - lambda nu, lengthscale, X: B.reshape( + def mean_variance(nu, lengthscale, X): + return B.reshape( B.mean( kernel.K_diag({"nu": nu, "lengthscale": lengthscale}, X), squeeze=False ), 1, - ), # the reshape shields from a bug in lab present at least up to version 1.6.6 + ) # 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["nu"], params["lengthscale"], X, diff --git a/tests/kernels/test_product_kernel.py b/tests/kernels/test_product_kernel.py index 7fb39ef9..da21a402 100644 --- a/tests/kernels/test_product_kernel.py +++ b/tests/kernels/test_product_kernel.py @@ -14,7 +14,7 @@ @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() + key = np.random.RandomState(0) key, xs_factor1 = factor1.random(key, 10) key, xs_factor2 = factor2.random(key, 10) @@ -25,8 +25,6 @@ def test_kernel_is_product_of_heat_kernels(factor1, factor2, nu, lengthscale, ba def K_diff(nu, lengthscale, xs_factor1, xs_factor2): params = {"nu": nu, "lengthscale": lengthscale} - # common_dtype = B.promote_dtypes(B.dtype(xs_factor1), B.dtype(xs_factor2)) - # xs_product = make_product([B.cast(common_dtype, xs_factor1), B.cast(common_dtype, xs_factor2)]) xs_product = make_product([xs_factor1, xs_factor2]) K_product = product_kernel.K(params, xs_product, xs_product) diff --git a/tests/sampling/test_samplers.py b/tests/sampling/test_samplers.py index f2493f54..da6eb970 100644 --- a/tests/sampling/test_samplers.py +++ b/tests/sampling/test_samplers.py @@ -17,16 +17,19 @@ def test_output_shape_and_backend(backend, feature_map_and_friends): params = kernel.init_params() sample_paths = sampler(feature_map, s=_NUM_SAMPLES) - key = np.random.RandomState() + key = np.random.RandomState(0) key, X = space.random(key, 50) + def sample(nu, lengthscale, X): + return sample_paths( + X, {"nu": nu, "lengthscale": lengthscale}, 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), - lambda nu, lengthscale, X: sample_paths( - X, {"nu": nu, "lengthscale": lengthscale}, key=create_random_state(backend) - )[1], + sample, params["nu"], params["lengthscale"], X, diff --git a/tests/spaces/test_circle.py b/tests/spaces/test_circle.py index dda7e95d..30efdfb7 100644 --- a/tests/spaces/test_circle.py +++ b/tests/spaces/test_circle.py @@ -39,14 +39,15 @@ def test_equivalence_kernel(nu, atol, backend): kernel = MaternGeometricKernel(Circle()) + def kern(nu, lengthscale, X, X2): + return kernel.K({"nu": nu, "lengthscale": lengthscale}, X, X2) + # Check that MaternGeometricKernel on Circle() coincides with the # periodic summation of the respective Euclidean Matérn kernel. check_function_with_backend( backend, result, - lambda nu, lengthscale, X, X2: kernel.K( - {"nu": nu, "lengthscale": lengthscale}, X, X2 - ), + kern, np.array([nu]), np.array([1.0]), inputs, diff --git a/tests/spaces/test_eigenfunctions.py b/tests/spaces/test_eigenfunctions.py index 4fc21053..c7c933c7 100644 --- a/tests/spaces/test_eigenfunctions.py +++ b/tests/spaces/test_eigenfunctions.py @@ -41,7 +41,7 @@ def inputs(request): num_levels = min(50, num_levels) eigenfunctions = space.get_eigenfunctions(num_levels) - key = np.random.RandomState() + 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) @@ -99,7 +99,7 @@ def test_orthonormality(inputs, backend): "CompactMatrixLieGroup subclasses do not currently support eigenfunction evaluation" ) - key = np.random.RandomState() + key = np.random.RandomState(0) key, xs = space.random(key, 10000) # Check that the eigenfunctions are orthonormal by comparing a Monte Carlo diff --git a/tests/spaces/test_hyperbolic.py b/tests/spaces/test_hyperbolic.py index c76a0d18..a11df283 100644 --- a/tests/spaces/test_hyperbolic.py +++ b/tests/spaces/test_hyperbolic.py @@ -18,7 +18,7 @@ def test_equivalence_kernel(dim, lengthscale, backend): space = Hyperbolic(dim) - key = np.random.RandomState() + key = np.random.RandomState(0) key, X = space.random(key, 6) X2 = X.copy() @@ -30,6 +30,9 @@ def test_equivalence_kernel(dim, lengthscale, backend): kernel = MaternGeometricKernel(space, key=create_random_state(backend)) + def kern(nu, lengthscale, X, X2): + return kernel.K({"nu": nu, "lengthscale": lengthscale}, X, X2) + # 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. @@ -38,9 +41,7 @@ def test_equivalence_kernel(dim, lengthscale, backend): check_function_with_backend( backend, result, - lambda nu, lengthscale, X, X2: kernel.K( - {"nu": nu, "lengthscale": lengthscale}, X, X2 - ), + kern, np.array([np.inf]), np.array([lengthscale]), X, diff --git a/tests/spaces/test_hypercube_graph.py b/tests/spaces/test_hypercube_graph.py index 745c654f..b41333cb 100644 --- a/tests/spaces/test_hypercube_graph.py +++ b/tests/spaces/test_hypercube_graph.py @@ -24,7 +24,7 @@ def inputs(request) -> Tuple[B.Numeric]: 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) @@ -55,15 +55,16 @@ def test_against_analytic_heat_kernel(inputs, lengthscale, backend): kernel = MaternGeometricKernel(space) + def kern(nu, lengthscale, X, X2): + return kernel.K({"nu": nu, "lengthscale": lengthscale}, X, X2) + # Check that MaternGeometricKernel on HypercubeGraph with nu=infinity # coincides with the closed form expression for the heat kernel on the # hypercube graph. check_function_with_backend( backend, result, - lambda nu, lengthscale, X, X2: kernel.K( - {"nu": nu, "lengthscale": lengthscale}, X, X2 - ), + kern, np.array([np.inf]), lengthscale, X, diff --git a/tests/spaces/test_lie_groups.py b/tests/spaces/test_lie_groups.py index f64b5a41..c5a1d327 100644 --- a/tests/spaces/test_lie_groups.py +++ b/tests/spaces/test_lie_groups.py @@ -25,7 +25,7 @@ def inputs(request): num_levels = min(10, default_num(space)) eigenfunctions = space.get_eigenfunctions(num_levels) - key = np.random.RandomState() + 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) @@ -110,7 +110,7 @@ def test_characters_orthogonal(inputs, backend): group, eigenfunctions, _, _ = inputs num_samples = 10000 - key = np.random.RandomState() + key = np.random.RandomState(0) _, X = group.random(key, num_samples) def all_char_vals(X): diff --git a/tests/spaces/test_product_discrete_spectrum_space.py b/tests/spaces/test_product_discrete_spectrum_space.py index 2aa2910b..5fada669 100644 --- a/tests/spaces/test_product_discrete_spectrum_space.py +++ b/tests/spaces/test_product_discrete_spectrum_space.py @@ -24,7 +24,7 @@ def test_heat_kernel_is_product_of_heat_kernels(factor1, factor2, lengthscale, b factor1, factor2, num_levels=_NUM_LEVELS**2, num_levels_per_space=_NUM_LEVELS ) - key = np.random.RandomState() + key = np.random.RandomState(0) key, xs_factor1 = factor1.random(key, 10) key, xs_factor2 = factor2.random(key, 10) diff --git a/tests/spaces/test_spd.py b/tests/spaces/test_spd.py index bf1327a3..f5f9e3d1 100644 --- a/tests/spaces/test_spd.py +++ b/tests/spaces/test_spd.py @@ -14,7 +14,7 @@ def test_equivalence_kernel(lengthscale, backend): space = SymmetricPositiveDefiniteMatrices(2) - key = np.random.RandomState() + key = np.random.RandomState(0) key, X = space.random(key, 5) X2 = X.copy() @@ -23,6 +23,9 @@ def test_equivalence_kernel(lengthscale, backend): kernel = MaternGeometricKernel(space, key=create_random_state(backend)) + def kern(nu, lengthscale, X, X2): + return kernel.K({"nu": nu, "lengthscale": lengthscale}, X, X2) + # 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 @@ -30,9 +33,7 @@ def test_equivalence_kernel(lengthscale, backend): check_function_with_backend( backend, result, - lambda nu, lengthscale, X, X2: kernel.K( - {"nu": nu, "lengthscale": lengthscale}, X, X2 - ), + kern, np.array([np.inf]), np.array([lengthscale]), X, diff --git a/tests/utils/test_kernel_formulas.py b/tests/utils/test_kernel_formulas.py index 8013f687..427acfc9 100644 --- a/tests/utils/test_kernel_formulas.py +++ b/tests/utils/test_kernel_formulas.py @@ -16,7 +16,7 @@ def test_hypercube_graph_heat_kernel(d, lengthscale, backend): space = HypercubeGraph(d) - 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) @@ -24,14 +24,17 @@ def test_hypercube_graph_heat_kernel(d, lengthscale, backend): 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, - lambda lengthscale, X, X2: hypercube_graph_heat_kernel( - lengthscale, X, X2, normalized_laplacian=False - ), + heat_kernel, np.array([lengthscale]), X, X2, @@ -58,9 +61,7 @@ def test_hypercube_graph_heat_kernel(d, lengthscale, backend): check_function_with_backend( backend, result, - lambda lengthscale, X, X2: hypercube_graph_heat_kernel( - lengthscale, X, X2, normalized_laplacian=False - ), + 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 index 56604570..08020cd9 100644 --- a/tests/utils/test_manifold_utils.py +++ b/tests/utils/test_manifold_utils.py @@ -12,7 +12,7 @@ def test_hyperboloid_distance(dim, backend): space = Hyperbolic(dim=dim) - key = np.random.RandomState() + 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) diff --git a/tests/utils/test_special_functions.py b/tests/utils/test_special_functions.py index d8e57a23..72f8f398 100644 --- a/tests/utils/test_special_functions.py +++ b/tests/utils/test_special_functions.py @@ -66,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, ) @@ -91,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, From 35ca2d975c5962e7c8b51df7cb409f84886cfb2c Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Mon, 25 Nov 2024 11:36:17 +0100 Subject: [PATCH 34/35] Incorporate the feedback from stoprightthere's review [continued] --- tests/kernels/test_matern_karhunenloeve_kernel.py | 2 +- tests/spaces/test_circle.py | 4 +++- tests/spaces/test_eigenfunctions.py | 2 +- tests/spaces/test_eigenvalues.py | 3 ++- tests/spaces/test_hyperbolic.py | 11 ++++++++--- tests/spaces/test_hypercube_graph.py | 2 +- tests/spaces/test_hypersphere.py | 2 +- tests/spaces/test_spd.py | 11 ++++++++--- 8 files changed, 25 insertions(+), 12 deletions(-) diff --git a/tests/kernels/test_matern_karhunenloeve_kernel.py b/tests/kernels/test_matern_karhunenloeve_kernel.py index 62a23e8a..f94e7013 100644 --- a/tests/kernels/test_matern_karhunenloeve_kernel.py +++ b/tests/kernels/test_matern_karhunenloeve_kernel.py @@ -71,7 +71,7 @@ def eigenvalues(nu, lengthscale): eigenvalues, params["nu"], params["lengthscale"], - compare_to_result=lambda res, f_out: f_out.shape == res, + compare_to_result=lambda res, f_out: B.shape(f_out) == res, ) diff --git a/tests/spaces/test_circle.py b/tests/spaces/test_circle.py index 30efdfb7..119c6b57 100644 --- a/tests/spaces/test_circle.py +++ b/tests/spaces/test_circle.py @@ -14,7 +14,9 @@ from ..helper import check_function_with_backend -@pytest.mark.parametrize("nu, atol", [(0.5, 1), (1.5, 3), (2.5, 5), (np.inf, 6)]) +@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: diff --git a/tests/spaces/test_eigenfunctions.py b/tests/spaces/test_eigenfunctions.py index c7c933c7..c6a98b93 100644 --- a/tests/spaces/test_eigenfunctions.py +++ b/tests/spaces/test_eigenfunctions.py @@ -48,7 +48,7 @@ def inputs(request): # These weights are used for testing the weighted outerproduct, they # should be positive. - weights = np.random.rand(eigenfunctions.num_levels, 1) ** 2 + 0.01 + weights = np.random.rand(eigenfunctions.num_levels, 1) ** 2 + 1e-5 return space, eigenfunctions, X, X2, weights diff --git a/tests/spaces/test_eigenvalues.py b/tests/spaces/test_eigenvalues.py index bccc02f2..6b04fb8b 100644 --- a/tests/spaces/test_eigenvalues.py +++ b/tests/spaces/test_eigenvalues.py @@ -22,7 +22,8 @@ def inputs(request): - 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). + - 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) diff --git a/tests/spaces/test_hyperbolic.py b/tests/spaces/test_hyperbolic.py index a11df283..6353b6dc 100644 --- a/tests/spaces/test_hyperbolic.py +++ b/tests/spaces/test_hyperbolic.py @@ -33,6 +33,13 @@ def test_equivalence_kernel(dim, lengthscale, backend): def kern(nu, lengthscale, X, X2): return kernel.K({"nu": nu, "lengthscale": lengthscale}, X, X2) + 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. @@ -46,7 +53,5 @@ def kern(nu, lengthscale, X, X2): np.array([lengthscale]), X, X2, - compare_to_result=lambda res, f_out: np.linalg.norm(res - B.to_numpy(f_out)) - / np.sqrt(res.shape[0] * res.shape[1]) - < 1e-1, + compare_to_result=compare_to_result, ) diff --git a/tests/spaces/test_hypercube_graph.py b/tests/spaces/test_hypercube_graph.py index b41333cb..cf489292 100644 --- a/tests/spaces/test_hypercube_graph.py +++ b/tests/spaces/test_hypercube_graph.py @@ -31,7 +31,7 @@ def inputs(request) -> Tuple[B.Numeric]: # These weights are used for testing the weighted outerproduct, they # should be positive. - weights = np.random.rand(eigenfunctions.num_levels, 1) ** 2 + 0.01 + weights = np.random.rand(eigenfunctions.num_levels, 1) ** 2 + 1e-5 return space, eigenfunctions, X, X2, weights diff --git a/tests/spaces/test_hypersphere.py b/tests/spaces/test_hypersphere.py index 7cfb20fe..006469ad 100644 --- a/tests/spaces/test_hypersphere.py +++ b/tests/spaces/test_hypersphere.py @@ -58,4 +58,4 @@ def heat_kernel(t, x, y): lapf = manifold_laplacian(x, hypersphere, egrad, ehess) # Check that they match - assert np.isclose(dfdt.detach().numpy(), lapf, atol=1.0e-3) + np.testing.assert_allclose(dfdt.detach().numpy(), lapf, atol=1e-3) diff --git a/tests/spaces/test_spd.py b/tests/spaces/test_spd.py index f5f9e3d1..2cfa3d85 100644 --- a/tests/spaces/test_spd.py +++ b/tests/spaces/test_spd.py @@ -26,6 +26,13 @@ def test_equivalence_kernel(lengthscale, backend): def kern(nu, lengthscale, X, X2): return kernel.K({"nu": nu, "lengthscale": lengthscale}, X, X2) + 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 @@ -38,7 +45,5 @@ def kern(nu, lengthscale, X, X2): np.array([lengthscale]), X, X2, - compare_to_result=lambda res, f_out: np.linalg.norm(res - B.to_numpy(f_out)) - / np.sqrt(res.shape[0] * res.shape[1]) - < 1e-1, + compare_to_result=compare_to_result, ) From 3115c5e128cb7be15ea5eb773945860cec99c9d5 Mon Sep 17 00:00:00 2001 From: Viacheslav Borovitskiy Date: Wed, 27 Nov 2024 18:40:50 +0100 Subject: [PATCH 35/35] Incorporate the feedback from stoprightthere's review [continued 2] --- tests/feature_maps/test_feature_maps.py | 7 +-- tests/helper.py | 31 ++++++++-- tests/kernels/test_feature_map_kernel.py | 36 +++--------- .../test_matern_karhunenloeve_kernel.py | 57 +++++-------------- tests/sampling/test_samplers.py | 9 +-- tests/spaces/test_circle.py | 8 +-- tests/spaces/test_hyperbolic.py | 8 +-- tests/spaces/test_hypercube_graph.py | 8 +-- tests/spaces/test_spd.py | 8 +-- 9 files changed, 63 insertions(+), 109 deletions(-) diff --git a/tests/feature_maps/test_feature_maps.py b/tests/feature_maps/test_feature_maps.py index 1bf0d235..db7f6c90 100644 --- a/tests/feature_maps/test_feature_maps.py +++ b/tests/feature_maps/test_feature_maps.py @@ -69,11 +69,8 @@ def diff_kern_mats(params, X): check_function_with_backend( backend, np.zeros((X.shape[0], X.shape[0])), - lambda nu, lengthscale, X: diff_kern_mats( - {"nu": nu, "lengthscale": lengthscale}, X - ), - params["nu"], - params["lengthscale"], + diff_kern_mats, + params, X, atol=0.1, ) diff --git a/tests/helper.py b/tests/helper.py index 5c28b940..9c3613ad 100644 --- a/tests/helper.py +++ b/tests/helper.py @@ -149,6 +149,27 @@ def array_type(backend: str): 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, @@ -184,14 +205,16 @@ 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) diff --git a/tests/kernels/test_feature_map_kernel.py b/tests/kernels/test_feature_map_kernel.py index 08fb8f81..d87a289a 100644 --- a/tests/kernels/test_feature_map_kernel.py +++ b/tests/kernels/test_feature_map_kernel.py @@ -63,16 +63,12 @@ def test_K(inputs, backend, normalize, kernel): _, _, _, X, X2 = inputs params = kernel.init_params() - def kern(nu, lengthscale, X, X2): - return kernel.K({"nu": nu, "lengthscale": lengthscale}, X, X2) - # 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]), - kern, - params["nu"], - params["lengthscale"], + kernel.K, + params, X, X2, compare_to_result=lambda res, f_out: B.shape(f_out) == res, @@ -85,18 +81,12 @@ def test_K_one_param(inputs, backend, normalize, kernel): _, _, _, X, _ = inputs params = kernel.init_params() - def diff(nu, lengthscale, X): - return kernel.K({"nu": nu, "lengthscale": lengthscale}, X) - kernel.K( - {"nu": nu, "lengthscale": lengthscale}, X, X - ) - # Check that kernel.K(X) coincides with kernel.K(X, X). check_function_with_backend( backend, np.zeros((X.shape[0], X.shape[0])), - diff, - params["nu"], - params["lengthscale"], + lambda params, X: kernel.K(params, X) - kernel.K(params, X, X), + params, X, ) @@ -107,18 +97,12 @@ def test_K_diag(inputs, backend, normalize, kernel): _, _, _, X, _ = inputs params = kernel.init_params() - def diff(nu, lengthscale, X): - return kernel.K_diag({"nu": nu, "lengthscale": lengthscale}, X) - B.diag( - kernel.K({"nu": nu, "lengthscale": lengthscale}, X) - ) - # Check that kernel.K_diag coincides with the diagonal of kernel.K. check_function_with_backend( backend, np.zeros((X.shape[0],)), - diff, - params["nu"], - params["lengthscale"], + lambda params, X: kernel.K_diag(params, X) - B.diag(kernel.K(params, X)), + params, X, ) @@ -129,15 +113,11 @@ def test_normalize(inputs, backend, kernel): params = kernel.init_params() - def kern_diag(nu, lengthscale, X): - return kernel.K_diag({"nu": nu, "lengthscale": lengthscale}, X) - # Check that the variance of the kernel is constant 1. check_function_with_backend( backend, np.ones((X.shape[0],)), - kern_diag, - params["nu"], - params["lengthscale"], + kernel.K_diag, + params, X, ) diff --git a/tests/kernels/test_matern_karhunenloeve_kernel.py b/tests/kernels/test_matern_karhunenloeve_kernel.py index f94e7013..802d6e61 100644 --- a/tests/kernels/test_matern_karhunenloeve_kernel.py +++ b/tests/kernels/test_matern_karhunenloeve_kernel.py @@ -61,16 +61,12 @@ def test_eigenvalues_shape(inputs, backend): _, num_levels, kernel, _, _ = inputs params = kernel.init_params() - def eigenvalues(nu, lengthscale): - return kernel.eigenvalues({"nu": nu, "lengthscale": lengthscale}) - # Check that the eigenvalues have appropriate shape. check_function_with_backend( backend, (num_levels, 1), - eigenvalues, - params["nu"], - params["lengthscale"], + kernel.eigenvalues, + params, compare_to_result=lambda res, f_out: B.shape(f_out) == res, ) @@ -80,16 +76,12 @@ def test_eigenvalues_positive(inputs, backend): _, _, kernel, _, _ = inputs params = kernel.init_params() - def eigenvalues(nu, lengthscale): - return kernel.eigenvalues({"nu": nu, "lengthscale": lengthscale}) - # Check that the eigenvalues are nonnegative. check_function_with_backend( backend, None, - eigenvalues, - params["nu"], - params["lengthscale"], + kernel.eigenvalues, + params, compare_to_result=lambda _, f_out: np.all(B.to_numpy(f_out) >= 0), ) @@ -99,16 +91,12 @@ def test_eigenvalues_ordered(inputs, backend): _, _, kernel, _, _ = inputs params = kernel.init_params() - def eigenvalues(nu, lengthscale): - return kernel.eigenvalues({"nu": nu, "lengthscale": lengthscale}) - # Check that the eigenvalues are sorted in descending order. check_function_with_backend( backend, None, - eigenvalues, - params["nu"], - params["lengthscale"], + kernel.eigenvalues, + params, compare_to_result=lambda _, f_out: np.all( B.to_numpy(f_out)[:-1] >= B.to_numpy(f_out)[1:] - _EPS ), @@ -124,17 +112,13 @@ def test_K(inputs, backend): assert result.shape == (X.shape[0], X2.shape[0]), "K has incorrect shape" - def kern(nu, lengthscale, X, X2): - return kernel.K({"nu": nu, "lengthscale": lengthscale}, X, X2) - if backend != "numpy": # Check that kernel.K computed using `backend` coincides with the numpy result. check_function_with_backend( backend, result, - kern, - params["nu"], - params["lengthscale"], + kernel.K, + params, X, X2, ) @@ -147,16 +131,12 @@ def test_K_one_param(inputs, backend): result = kernel.K(params, X, X) - def kern(nu, lengthscale, X): - return kernel.K({"nu": nu, "lengthscale": lengthscale}, X) - # Check that kernel.K(X) coincides with kernel.K(X, X). check_function_with_backend( backend, result, - kern, - params["nu"], - params["lengthscale"], + kernel.K, + params, X, atol=1e-2 if isinstance(space, Mesh) else _EPS, ) @@ -171,16 +151,12 @@ def test_K_diag(inputs, backend): assert result.shape == (X.shape[0],), "The diagonal has incorrect shape" - def kern_diag(nu, lengthscale, X): - return kernel.K_diag({"nu": nu, "lengthscale": lengthscale}, X) - # Check that kernel.K_diag coincides with the diagonal of kernel.K. check_function_with_backend( backend, result, - kern_diag, - params["nu"], - params["lengthscale"], + kernel.K_diag, + params, X, atol=1e-2 if isinstance(space, Mesh) else _EPS, ) @@ -199,11 +175,9 @@ def test_normalize(inputs, backend): key, 1000 ) # we need a large sample to get a good estimate of the mean variance - def mean_variance(nu, lengthscale, X): + def mean_variance(params, X): return B.reshape( - B.mean( - kernel.K_diag({"nu": nu, "lengthscale": lengthscale}, X), squeeze=False - ), + 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 @@ -212,8 +186,7 @@ def mean_variance(nu, lengthscale, X): backend, np.array([1.0]), mean_variance, - params["nu"], - params["lengthscale"], + params, X, atol=0.2, # very loose, but helps make sure the result is close to 1 ) diff --git a/tests/sampling/test_samplers.py b/tests/sampling/test_samplers.py index da6eb970..557654d0 100644 --- a/tests/sampling/test_samplers.py +++ b/tests/sampling/test_samplers.py @@ -20,18 +20,15 @@ def test_output_shape_and_backend(backend, feature_map_and_friends): key = np.random.RandomState(0) key, X = space.random(key, 50) - def sample(nu, lengthscale, X): - return sample_paths( - X, {"nu": nu, "lengthscale": lengthscale}, key=create_random_state(backend) - )[1] + 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["nu"], - params["lengthscale"], + params, X, compare_to_result=lambda res, f_out: B.shape(f_out) == res, ) diff --git a/tests/spaces/test_circle.py b/tests/spaces/test_circle.py index 119c6b57..c7e59d7a 100644 --- a/tests/spaces/test_circle.py +++ b/tests/spaces/test_circle.py @@ -41,17 +41,13 @@ def test_equivalence_kernel(nu, atol, backend): kernel = MaternGeometricKernel(Circle()) - def kern(nu, lengthscale, X, X2): - return kernel.K({"nu": nu, "lengthscale": lengthscale}, X, X2) - # Check that MaternGeometricKernel on Circle() coincides with the # periodic summation of the respective Euclidean Matérn kernel. check_function_with_backend( backend, result, - kern, - np.array([nu]), - np.array([1.0]), + kernel.K, + {"nu": np.array([nu]), "lengthscale": np.array([1.0])}, inputs, inputs2, atol=atol, diff --git a/tests/spaces/test_hyperbolic.py b/tests/spaces/test_hyperbolic.py index 6353b6dc..9b51260f 100644 --- a/tests/spaces/test_hyperbolic.py +++ b/tests/spaces/test_hyperbolic.py @@ -30,9 +30,6 @@ def test_equivalence_kernel(dim, lengthscale, backend): kernel = MaternGeometricKernel(space, key=create_random_state(backend)) - def kern(nu, lengthscale, X, X2): - return kernel.K({"nu": nu, "lengthscale": lengthscale}, X, X2) - def compare_to_result(res, f_out): return ( np.linalg.norm(res - B.to_numpy(f_out)) @@ -48,9 +45,8 @@ def compare_to_result(res, f_out): check_function_with_backend( backend, result, - kern, - np.array([np.inf]), - np.array([lengthscale]), + 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 cf489292..cb5ac9c9 100644 --- a/tests/spaces/test_hypercube_graph.py +++ b/tests/spaces/test_hypercube_graph.py @@ -55,18 +55,14 @@ def test_against_analytic_heat_kernel(inputs, lengthscale, backend): kernel = MaternGeometricKernel(space) - def kern(nu, lengthscale, X, X2): - return kernel.K({"nu": nu, "lengthscale": lengthscale}, X, X2) - # Check that MaternGeometricKernel on HypercubeGraph with nu=infinity # coincides with the closed form expression for the heat kernel on the # hypercube graph. check_function_with_backend( backend, result, - kern, - 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_spd.py b/tests/spaces/test_spd.py index 2cfa3d85..9e4729e2 100644 --- a/tests/spaces/test_spd.py +++ b/tests/spaces/test_spd.py @@ -23,9 +23,6 @@ def test_equivalence_kernel(lengthscale, backend): kernel = MaternGeometricKernel(space, key=create_random_state(backend)) - def kern(nu, lengthscale, X, X2): - return kernel.K({"nu": nu, "lengthscale": lengthscale}, X, X2) - def compare_to_result(res, f_out): return ( np.linalg.norm(res - B.to_numpy(f_out)) @@ -40,9 +37,8 @@ def compare_to_result(res, f_out): check_function_with_backend( backend, result, - kern, - np.array([np.inf]), - np.array([lengthscale]), + kernel.K, + {"nu": np.array([np.inf]), "lengthscale": np.array([lengthscale])}, X, X2, compare_to_result=compare_to_result,