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

1import lab as B 

2import numpy as np 

3import pytest 

4 

5from geometric_kernels.kernels import MaternGeometricKernel 

6from geometric_kernels.spaces import SymmetricPositiveDefiniteMatrices 

7from geometric_kernels.utils.kernel_formulas import spd_heat_kernel_2x2 

8 

9from ..helper import check_function_with_backend, create_random_state 

10 

11 

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

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

14def test_equivalence_kernel(lengthscale, backend): 

15 space = SymmetricPositiveDefiniteMatrices(2) 

16 

17 key = np.random.RandomState(0) 

18 key, X = space.random(key, 5) 

19 X2 = X.copy() 

20 

21 t = lengthscale * lengthscale / 2 

22 result = spd_heat_kernel_2x2(t, X, X2) 

23 

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

25 

26 def compare_to_result(res, f_out): 

27 return ( 

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

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

30 < 1e-1 

31 ) 

32 

33 # Check that MaternGeometricKernel on SymmetricPositiveDefiniteMatrices(2) 

34 # with nu=inf coincides with the semi-analytic formula from :cite:t:`sawyer1992`. 

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

36 # the two covariance matrices. 

37 check_function_with_backend( 

38 backend, 

39 result, 

40 kernel.K, 

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

42 X, 

43 X2, 

44 compare_to_result=compare_to_result, 

45 )