Coverage for tests/feature_maps/test_feature_maps.py: 100%

1import lab as B 

2import numpy as np 

3import pytest 

4 

5from geometric_kernels.feature_maps import RandomPhaseFeatureMapCompact 

6from geometric_kernels.kernels import MaternGeometricKernel, default_feature_map 

7from geometric_kernels.kernels.matern_kernel import default_num 

8from geometric_kernels.spaces import NoncompactSymmetricSpace 

9from geometric_kernels.utils.utils import make_deterministic 

10 

11from ..helper import check_function_with_backend, create_random_state, spaces 

12 

13 

14@pytest.fixture( 

15 params=spaces(), 

16 ids=str, 

17) 

18def feature_map_and_friends(request, backend): 

19 """ 

20 Returns a tuple (feature_map, kernel, space) where: 

21 - feature_map is the `default_feature_map` of the `kernel`, 

22 - kernel is the `MaternGeometricKernel` on the `space`, with a reasonably 

23 small value of `num`, 

24 - space = request.param, 

25 

26 `backend` parameter is required to create a random state for the feature 

27 map, if it requires one. 

28 """ 

29 space = request.param 

30 

31 if isinstance(space, NoncompactSymmetricSpace): 

32 kernel = MaternGeometricKernel( 

33 space, key=create_random_state(backend), num=min(default_num(space), 100) 

34 ) 

35 else: 

36 kernel = MaternGeometricKernel(space, num=min(default_num(space), 3)) 

37 

38 feature_map = default_feature_map(kernel=kernel) 

39 if isinstance(feature_map, RandomPhaseFeatureMapCompact): 

40 # RandomPhaseFeatureMapCompact requires a key. Note: normally, 

41 # RandomPhaseFeatureMapNoncompact, RejectionSamplingFeatureMapHyperbolic, 

42 # and RejectionSamplingFeatureMapSPD also require a key, but when they 

43 # are obtained from an already constructed kernel's feature map, the key 

44 # is already provided and fixed in the similar way as we do just below. 

45 feature_map = make_deterministic(feature_map, key=create_random_state(backend)) 

46 

47 return feature_map, kernel, space 

48 

49 

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

51def test_feature_map_approximates_kernel(backend, feature_map_and_friends): 

52 feature_map, kernel, space = feature_map_and_friends 

53 

54 params = kernel.init_params() 

55 

56 key = np.random.RandomState(0) 

57 key, X = space.random(key, 50) 

58 

59 def diff_kern_mats(params, X): 

60 _, embedding = feature_map(X, params) 

61 

62 kernel_mat = kernel.K(params, X, X) 

63 kernel_mat_alt = B.matmul(embedding, B.T(embedding)) 

64 

65 return kernel_mat - kernel_mat_alt 

66 

67 # Check that, approximately, k(X, X) = <phi(X), phi(X)>, where k is the 

68 # kernel and phi is the feature map. 

69 check_function_with_backend( 

70 backend, 

71 np.zeros((X.shape[0], X.shape[0])), 

72 diff_kern_mats, 

73 params, 

74 X, 

75 atol=0.1, 

76 )