Coverage for tests/spaces/test_hyperbolic.py: 100%

1import lab as B 

2import numpy as np 

3import pytest 

4 

5from geometric_kernels.kernels import MaternGeometricKernel 

6from geometric_kernels.spaces import Hyperbolic 

7from geometric_kernels.utils.kernel_formulas import ( 

8 hyperbolic_heat_kernel_even, 

9 hyperbolic_heat_kernel_odd, 

10) 

11 

12from ..helper import check_function_with_backend, create_random_state 

13 

14 

15@pytest.mark.parametrize("dim", [2, 3, 5, 7]) 

16@pytest.mark.parametrize("lengthscale", [2.0]) 

17@pytest.mark.parametrize("backend", ["numpy", "tensorflow", "torch", "jax"]) 

18def test_equivalence_kernel(dim, lengthscale, backend): 

19 space = Hyperbolic(dim) 

20 

21 key = np.random.RandomState(0) 

22 key, X = space.random(key, 6) 

23 X2 = X.copy() 

24 

25 t = lengthscale * lengthscale / 2 

26 if dim % 2 == 1: 

27 result = hyperbolic_heat_kernel_odd(dim, t, X, X2) 

28 else: 

29 result = hyperbolic_heat_kernel_even(dim, t, X, X2) 

30 

31 kernel = MaternGeometricKernel(space, key=create_random_state(backend)) 

32 

33 def compare_to_result(res, f_out): 

34 return ( 

35 np.linalg.norm(res - B.to_numpy(f_out)) 

36 / np.sqrt(res.shape[0] * res.shape[1]) 

37 < 1e-1 

38 ) 

39 

40 # Check that MaternGeometricKernel on Hyperbolic(dim) with nu=inf coincides 

41 # with the well-known analytic formula for the heat kernel on the hyperbolic 

42 # space in odd dimensions and semi-analytic formula in even dimensions. 

43 # We are checking the equivalence on average, computing the norm between 

44 # the two covariance matrices. 

45 check_function_with_backend( 

46 backend, 

47 result, 

48 kernel.K, 

49 {"nu": np.array([np.inf]), "lengthscale": np.array([lengthscale])}, 

50 X, 

51 X2, 

52 compare_to_result=compare_to_result, 

53 )