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

1import lab as B 

2import numpy as np 

3import pytest 

4 

5from geometric_kernels.kernels import MaternGeometricKernel 

6from geometric_kernels.spaces.circle import Circle 

7from geometric_kernels.utils.kernel_formulas import ( 

8 euclidean_matern_12_kernel, 

9 euclidean_matern_32_kernel, 

10 euclidean_matern_52_kernel, 

11 euclidean_rbf_kernel, 

12) 

13 

14from ..helper import check_function_with_backend 

15 

16 

17@pytest.mark.parametrize( 

18 "nu, atol", [(0.5, 1e-1), (1.5, 1e-3), (2.5, 1e-3), (np.inf, 1e-5)] 

19) 

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

21def test_equivalence_kernel(nu, atol, backend): 

22 if nu == 0.5: 

23 analytic_kernel = euclidean_matern_12_kernel 

24 elif nu == 1.5: 

25 analytic_kernel = euclidean_matern_32_kernel 

26 elif nu == 2.5: 

27 analytic_kernel = euclidean_matern_52_kernel 

28 elif nu == np.inf: 

29 analytic_kernel = euclidean_rbf_kernel 

30 

31 inputs = np.random.uniform(0, 2 * np.pi, size=(5, 1)) 

32 inputs2 = np.random.uniform(0, 2 * np.pi, size=(3, 1)) 

33 

34 # Compute kernel using periodic summation 

35 geodesic = inputs[:, None, :] - inputs2[None, :, :] # [N, N2, 1] 

36 all_distances = ( 

37 geodesic + np.array([i * 2 * np.pi for i in range(-10, 10)])[None, None, :] 

38 ) 

39 all_distances = B.abs(all_distances) 

40 result = B.to_numpy(B.sum(analytic_kernel(all_distances), axis=2)) 

41 

42 kernel = MaternGeometricKernel(Circle()) 

43 

44 # Check that MaternGeometricKernel on Circle() coincides with the 

45 # periodic summation of the respective Euclidean Matérn kernel. 

46 check_function_with_backend( 

47 backend, 

48 result, 

49 kernel.K, 

50 {"nu": np.array([nu]), "lengthscale": np.array([1.0])}, 

51 inputs, 

52 inputs2, 

53 atol=atol, 

54 )