Coverage for geometric_kernels/utils/kernel_formulas/spd.py: 90%

1""" 

2Implements an alternative formula for the heat kernel on the manifold of 

3symmetric positive definite matrices by :cite:t:`sawyer1992`. 

4 

5The implementation is adapted from https://github.com/imbirik/LieStationaryKernels. 

6Since the resulting approximation 

7* can fail to be positive semi-definite, 

8* is very slow, 

9* and is rather numerically unstable, 

10it is not recommended to use it in practice. The implementation is provided 

11mainly for testing purposes. 

12""" 

13 

14import lab as B 

15import numpy as np 

16import scipy 

17from beartype.typing import Optional 

18 

19 

20def _spd_heat_kernel_2x2_base( 

21 t: float, 

22 x: B.NPNumeric, 

23 x2: Optional[B.NPNumeric] = None, 

24) -> float: 

25 """ 

26 The semi-analytic formula for the heat kernel on manifold of symmetric 

27 positive definite matrices 2x2 from :cite:t:`sawyer1992`. The implementation 

28 is adapted from https://github.com/imbirik/LieStationaryKernels. 

29 

30 :param t: 

31 The time parameter, a positive float. 

32 :param x: 

33 A single input, an array of shape [2, 2]. 

34 :param x2: 

35 A single input, an array of shape [2, 2]. If None, defaults to x. 

36 

37 :return: 

38 An approximation of the kernel value k(x, x2), a float. The kernel is not 

39 normalized, i.e. k(x, x) may be an arbitrary (implementation-dependent) 

40 positive number. For the normalized kernel which can also handle batch 

41 inputs outputting covariance matrices, use :func:`sawyer_heat_kernel`. 

42 """ 

43 if x2 is None: 

44 x2 = x 

45 

46 if B.shape(x) != (2, 2): 

47 raise ValueError("`x` must have shape [2, 2].") 

48 if x2.shape != (2, 2): 

49 raise ValueError("`x2` must have shape [2, 2].") 

50 

51 cl_1 = np.linalg.cholesky(x) 

52 cl_2 = np.linalg.cholesky(x2) 

53 diff = np.linalg.inv(cl_2) @ cl_1 

54 _, singular_values, _ = np.linalg.svd(diff) 

55 # Note: singular values that np.linalg.svd outputs are sorted, the following 

56 # code relies on this fact. 

57 H1, H2 = np.log(singular_values[0]), np.log(singular_values[1]) 

58 if H1 < H2: 

59 raise RuntimeError("Expected `np.linalg.svd` to return sorted eigenvalues.") 

60 

61 r_H_sq = H1 * H1 + H2 * H2 

62 alpha = H1 - H2 

63 

64 # Non-integral part 

65 result = 1.0 

66 result *= np.exp(-r_H_sq / (4 * t)) 

67 

68 # Integrand 

69 def link_function(x): 

70 if x < 1e-5: 

71 x = 1e-5 

72 res = 1.0 

73 res *= 2 * x + alpha 

74 res *= np.exp(-x * (x + alpha) / (2 * t)) 

75 res *= pow(np.sinh(x) * np.sinh(x + alpha), -1 / 2) 

76 return res 

77 

78 # Evaluating the integral 

79 

80 # scipy.integrate.quad is much more accurate than np.trapz with 

81 # b_vals = np.logspace(-3., 1, 1000), at least if we believe 

82 # that Mathematica's NIntegrate is accurate. Also, you might think that 

83 # scipy.integrate.quad_vec can be used to compute a whole covariance matrix 

84 # at once. However, it seems to show terrible accuracy in this case. 

85 

86 integral, error = scipy.integrate.quad(link_function, 0, np.inf) 

87 

88 result *= integral 

89 

90 return result 

91 

92 

93def spd_heat_kernel_2x2( 

94 t: float, 

95 X: B.NPNumeric, 

96 X2: Optional[B.NPNumeric] = None, 

97) -> B.NPNumeric: 

98 """ 

99 The semi-analytic formula for the heat kernel on manifold of symmetric 

100 positive definite matrices 2x2 from :cite:t:`sawyer1992`, normalized to 

101 have k(x, x) = 1 for all x. The implementation is adapted from 

102 https://github.com/imbirik/LieStationaryKernels. 

103 

104 :param t: 

105 The time parameter, a positive float. 

106 :param X: 

107 A batch of inputs, an array of shape [N, 2, 2]. 

108 :param X2: 

109 A batch of inputs, an array of shape [N2, 2, 2]. If None, defaults to X. 

110 

111 :return: 

112 The kernel matrix, an array of shape [N, N2]. The kernel is normalized, 

113 i.e. k(x, x) = 1 for all x. 

114 """ 

115 

116 if X2 is None: 

117 X2 = X 

118 

119 normalization = _spd_heat_kernel_2x2_base(t, np.eye(2, 2)) 

120 

121 result = np.zeros((X.shape[0], X2.shape[0])) 

122 

123 # This is a very inefficient implementation, but it will do for tests. The 

124 # straightforward vectorization of _sawyer_heat_kernel_base is not possible 

125 # due to scipy.integrate.quad_vec giving very bad accuracy in this case. 

126 for i, x in enumerate(X): 

127 for j, x2 in enumerate(X2): 

128 result[i, j] = _spd_heat_kernel_2x2_base(t, x, x2) / normalization 

129 

130 return result