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

1import warnings 

2 

3import lab as B 

4import numpy as np 

5import pytest 

6 

7from geometric_kernels.jax import * # noqa 

8from geometric_kernels.kernels import MaternGeometricKernel 

9from geometric_kernels.spaces import Graph 

10from geometric_kernels.tensorflow import * # noqa 

11from geometric_kernels.torch import * # noqa 

12 

13from ..data import ( 

14 TEST_GRAPH_ADJACENCY, 

15 TEST_GRAPH_LAPLACIAN, 

16 TEST_GRAPH_LAPLACIAN_NORMALIZED, 

17) 

18from ..helper import check_function_with_backend, np_to_backend 

19 

20warnings.filterwarnings("ignore", category=RuntimeWarning, module="scipy") 

21 

22A = TEST_GRAPH_ADJACENCY 

23L = TEST_GRAPH_LAPLACIAN 

24 

25 

26@pytest.mark.parametrize("normalized", [True, False]) 

27@pytest.mark.parametrize( 

28 "backend", ["numpy", "tensorflow", "torch", "jax", "scipy_sparse"] 

29) 

30def test_laplacian(normalized, backend): 

31 

32 # Check that the Laplacian is computed correctly. 

33 check_function_with_backend( 

34 backend, 

35 TEST_GRAPH_LAPLACIAN if not normalized else TEST_GRAPH_LAPLACIAN_NORMALIZED, 

36 lambda adj: Graph(adj, normalize_laplacian=normalized)._laplacian, 

37 TEST_GRAPH_ADJACENCY, 

38 ) 

39 

40 

41@pytest.mark.parametrize( 

42 "L", [TEST_GRAPH_ADJACENCY.shape[0], TEST_GRAPH_ADJACENCY.shape[0] // 2] 

43) 

44@pytest.mark.parametrize("normalized", [True, False]) 

45@pytest.mark.parametrize( 

46 "backend", ["numpy", "tensorflow", "torch", "jax", "scipy_sparse"] 

47) 

48def test_eigendecomposition(L, normalized, backend): 

49 laplacian = np_to_backend( 

50 TEST_GRAPH_LAPLACIAN if not normalized else TEST_GRAPH_LAPLACIAN_NORMALIZED, 

51 backend, 

52 ) 

53 

54 def eigendiff(adj): 

55 graph = Graph(adj, normalize_laplacian=normalized) 

56 

57 eigenvalue_mat = B.diag_construct(graph.get_eigenvalues(L)[:, 0]) 

58 eigenvectors = graph.get_eigenvectors(L) 

59 # If the backend is scipy_sparse, convert eigenvalues/eigenvectors, 

60 # which are always supposed to be dense arrays, to sparse arrays. 

61 if backend == "scipy_sparse": 

62 import scipy.sparse as sp 

63 

64 eigenvalue_mat = sp.csr_array(eigenvalue_mat) 

65 eigenvectors = sp.csr_array(eigenvectors) 

66 

67 laplace_x_eigvecs = laplacian @ eigenvectors 

68 eigvals_x_eigvecs = eigenvectors @ eigenvalue_mat 

69 return laplace_x_eigvecs - eigvals_x_eigvecs 

70 

71 # Check that the computed eigenpairs are actually eigenpairs 

72 check_function_with_backend( 

73 backend, 

74 np.zeros((TEST_GRAPH_ADJACENCY.shape[0], L)), 

75 eigendiff, 

76 TEST_GRAPH_ADJACENCY, 

77 ) 

78 

79 

80@pytest.mark.parametrize("nu, lengthscale", [(1.0, 1.0), (2.0, 1.0), (np.inf, 1.0)]) 

81@pytest.mark.parametrize("sparse_adj", [True, False]) 

82@pytest.mark.parametrize("normalized", [True, False]) 

83@pytest.mark.parametrize( 

84 "backend", ["numpy", "tensorflow", "torch", "jax"] 

85) # 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. 

86def test_matern_kernels(nu, lengthscale, sparse_adj, normalized, backend): 

87 

88 laplacian = ( 

89 TEST_GRAPH_LAPLACIAN if not normalized else TEST_GRAPH_LAPLACIAN_NORMALIZED 

90 ) 

91 

92 evals_np, evecs_np = np.linalg.eigh(laplacian) 

93 evecs_np *= np.sqrt(laplacian.shape[0]) 

94 

95 def evaluate_kernel(adj, nu, lengthscale): 

96 dtype = B.dtype(adj) 

97 if sparse_adj: 

98 adj = np_to_backend(B.to_numpy(adj), "scipy_sparse") 

99 graph = Graph(adj, normalize_laplacian=normalized) 

100 kernel = MaternGeometricKernel(graph) 

101 return kernel.K( 

102 {"nu": nu, "lengthscale": lengthscale}, 

103 B.range(dtype, adj.shape[0])[:, None], 

104 ) 

105 

106 if nu < np.inf: 

107 K = ( 

108 evecs_np 

109 @ np.diag(np.power(evals_np + 2 * nu / lengthscale**2, -nu)) 

110 @ evecs_np.T 

111 ) 

112 else: 

113 K = evecs_np @ np.diag(np.exp(-(lengthscale**2) / 2 * evals_np)) @ evecs_np.T 

114 K = K / np.mean(K.diagonal()) 

115 

116 # Check that the kernel matrix is correctly computed 

117 check_function_with_backend( 

118 backend, 

119 K, 

120 evaluate_kernel, 

121 TEST_GRAPH_ADJACENCY, 

122 np.array([nu]), 

123 np.array([lengthscale]), 

124 )