Coverage for geometric_kernels/kernels/karhunen_loeve.py: 95%

1""" 

2This module provides the :class:`MaternKarhunenLoeveKernel` kernel, the basic 

3kernel for discrete spectrum spaces, subclasses of :class:`DiscreteSpectrumSpace`. 

4""" 

5 

6import lab as B 

7import numpy as np 

8from beartype.typing import Dict, Optional 

9 

10from geometric_kernels.kernels.base import BaseGeometricKernel 

11from geometric_kernels.lab_extras import from_numpy, is_complex 

12from geometric_kernels.spaces import DiscreteSpectrumSpace 

13from geometric_kernels.spaces.eigenfunctions import Eigenfunctions 

14from geometric_kernels.utils.utils import _check_1_vector, _check_field_in_params 

15 

16 

17class MaternKarhunenLoeveKernel(BaseGeometricKernel): 

18 r""" 

19 This class approximates Matérn kernel by its truncated Mercer decomposition, 

20 in terms of the eigenfunctions & eigenvalues of the Laplacian on the space. 

21 

22 .. math:: k(x, x') = \sum_{l=0}^{L-1} S(\sqrt\lambda_l) \sum_{s=1}^{d_l} f_{ls}(x) f_{ls}(x'), 

23 

24 where $\lambda_l$ and $f_{ls}(\cdot)$ are the eigenvalues and 

25 eigenfunctions of the Laplacian such that 

26 $\Delta f_{ls} = \lambda_l f_{ls}$, and $S(\cdot)$ is the spectrum 

27 of the Matérn kernel. The eigenvalues and eigenfunctions belong to the 

28 :class:`~.spaces.DiscreteSpectrumSpace` instance. 

29 

30 We denote 

31 

32 .. math:: G_l(\cdot, \cdot') = \sum_{s=1}^{d_l} f_{ls}(\cdot) f_{ls}(\cdot') 

33 

34 and term the sets $[f_{ls}]_{s=1}^{d_l}$ *"levels"*. 

35 

36 For many spaces, like the sphere, we can employ addition 

37 theorems to efficiently compute $G_l(\cdot, \cdot')$ without calculating 

38 the individual $f_{ls}(\cdot)$. Note that $\lambda_l$ are not required to 

39 be unique: it is possible that for some $l,l'$, $\lambda_l = \lambda_{l'}$. 

40 In other words, the "levels" do not necessarily correspond to full 

41 eigenspaces. A level may even correspond to a single eigenfunction. 

42 

43 .. note:: 

44 A brief introduction into the theory behind 

45 :class:`MaternKarhunenLoeveKernel` can be found in 

46 :doc:`this </theory/compact>` & :doc:`this </theory/addition_theorem>` 

47 documentation pages. 

48 

49 :param space: 

50 The space to define the kernel upon. 

51 :param num_levels: 

52 Number of levels to include in the summation. 

53 :param normalize: 

54 Whether to normalize kernel to have unit average variance. 

55 :param eigenvalues_laplacian: 

56 Allowing to pass the eigenvalues of the Laplacian directly, instead of 

57 computing them from the space. If provided, `eigenfunctions` must also 

58 be provided. Used for :class:`~.spaces.HodgeDiscreteSpectrumSpace`. 

59 :param eigenfunctions: 

60 Allowing to pass the eigenfunctions directly, instead of computing them 

61 from the space. If provided, `eigenvalues_laplacian` must also be provided. 

62 Used for :class:`~.spaces.HodgeDiscreteSpectrumSpace`. 

63 """ 

64 

65 def __init__( 

66 self, 

67 space: DiscreteSpectrumSpace, 

68 num_levels: int, 

69 normalize: bool = True, 

70 eigenvalues_laplacian: Optional[B.Numeric] = None, 

71 eigenfunctions: Optional[Eigenfunctions] = None, 

72 ): 

73 super().__init__(space) 

74 self.num_levels = num_levels # in code referred to as `L`. 

75 

76 if eigenvalues_laplacian is None: 

77 if eigenfunctions is not None: 

78 raise ValueError( 

79 "You must either provide both `eigenfunctions` and `eigenvalues_laplacian`, or none of the two." 

80 ) 

81 eigenvalues_laplacian = self.space.get_eigenvalues(self.num_levels) 

82 eigenfunctions = self.space.get_eigenfunctions(self.num_levels) 

83 else: 

84 if eigenfunctions is None: 

85 raise ValueError( 

86 "You must either provide both `eigenfunctions` and `eigenvalues_laplacian`, or none of the two." 

87 ) 

88 if eigenvalues_laplacian.shape != (num_levels, 1): 

89 raise ValueError( 

90 f"Expected `eigenvalues_laplacian` to have shape [num_levels={num_levels}, 1] but got {eigenvalues_laplacian.shape}" 

91 ) 

92 if eigenfunctions.num_levels != num_levels: 

93 raise ValueError( 

94 f"`num_levels` must coincide with `num_levels` in the provided `eigenfunctions`," 

95 f"but `num_levels`={num_levels} and `eigenfunctions.num_levels`={eigenfunctions.num_levels}" 

96 ) 

97 

98 self._eigenvalues_laplacian = eigenvalues_laplacian 

99 self._eigenfunctions = eigenfunctions 

100 self.normalize = normalize 

101 

102 @property 

103 def space(self) -> DiscreteSpectrumSpace: 

104 """ 

105 The space on which the kernel is defined. 

106 """ 

107 self._space: DiscreteSpectrumSpace 

108 return self._space 

109 

110 def init_params(self) -> Dict[str, B.NPNumeric]: 

111 """ 

112 Initializes the dict of the trainable parameters of the kernel. 

113 

114 Returns `dict(nu=np.array([np.inf]), lengthscale=np.array([1.0]))`. 

115 

116 This dict can be modified and is passed around into such methods as 

117 :meth:`~.K` or :meth:`~.K_diag`, as the `params` argument. 

118 

119 .. note:: 

120 The values in the returned dict are always of the NumPy array type. 

121 Thus, if you want to use some other backend for internal 

122 computations when calling :meth:`~.K` or :meth:`~.K_diag`, you 

123 need to replace the values with the analogs typed as arrays of 

124 the desired backend. 

125 """ 

126 params = dict(nu=np.array([np.inf]), lengthscale=np.array([1.0])) 

127 

128 return params 

129 

130 @staticmethod 

131 def spectrum( 

132 s: B.Numeric, nu: B.Numeric, lengthscale: B.Numeric, dimension: int 

133 ) -> B.Numeric: 

134 """ 

135 Static method computing the spectrum of the Matérn kernel with 

136 hyperparameters `nu` and `lengthscale` on the space with eigenvalues `s` 

137 and dimension `dimension`. 

138 

139 :param s: 

140 The eigenvalues of the space. 

141 :param nu: 

142 The smoothness parameter of the kernel. 

143 :param lengthscale: 

144 The length scale parameter of the kernel. 

145 :param dimension: 

146 The dimension of the space. 

147 

148 :return: 

149 The spectrum of the Matérn kernel. 

150 """ 

151 _check_1_vector(lengthscale, "lengthscale") 

152 _check_1_vector(nu, "nu") 

153 

154 # Note: 1.0 in safe_nu can be replaced by any finite positive value 

155 safe_nu = B.where(nu == np.inf, B.cast(B.dtype(lengthscale), np.r_[1.0]), nu) 

156 

157 # for nu == np.inf 

158 spectral_values_nu_infinite = B.exp( 

159 -(lengthscale**2) / 2.0 * B.cast(B.dtype(lengthscale), s) 

160 ) 

161 

162 # for nu < np.inf 

163 power = -safe_nu - dimension / 2.0 

164 base = 2.0 * safe_nu / lengthscale**2 + B.cast(B.dtype(safe_nu), s) 

165 spectral_values_nu_finite = base**power 

166 

167 return B.where( 

168 nu == np.inf, spectral_values_nu_infinite, spectral_values_nu_finite 

169 ) 

170 

171 @property 

172 def eigenfunctions(self) -> Eigenfunctions: 

173 """ 

174 Eigenfunctions of the kernel. 

175 """ 

176 return self._eigenfunctions 

177 

178 @property 

179 def eigenvalues_laplacian(self) -> B.Numeric: 

180 """ 

181 Eigenvalues of the Laplacian. 

182 """ 

183 return self._eigenvalues_laplacian 

184 

185 def eigenvalues(self, params: Dict[str, B.Numeric]) -> B.Numeric: 

186 """ 

187 Eigenvalues of the kernel. 

188 

189 :param params: 

190 Parameters of the kernel. Must contain keys `"lengthscale"` and 

191 `"nu"`. The shapes of `params["lengthscale"]` and `params["nu"]` 

192 are `(1,)`. 

193 

194 :return: 

195 An [L, 1]-shaped array. 

196 """ 

197 _check_field_in_params(params, "lengthscale") 

198 _check_1_vector(params["lengthscale"], 'params["lengthscale"]') 

199 

200 _check_field_in_params(params, "nu") 

201 _check_1_vector(params["nu"], 'params["nu"]') 

202 

203 spectral_values = self.spectrum( 

204 self.eigenvalues_laplacian, 

205 nu=params["nu"], 

206 lengthscale=params["lengthscale"], 

207 dimension=self.space.dimension, 

208 ) 

209 

210 if self.normalize: 

211 normalizer = B.sum( 

212 spectral_values 

213 * B.cast( 

214 B.dtype(spectral_values), 

215 from_numpy( 

216 spectral_values, 

217 self.eigenfunctions.num_eigenfunctions_per_level, 

218 )[:, None], 

219 ) 

220 ) 

221 return spectral_values / normalizer 

222 

223 return spectral_values 

224 

225 def K( 

226 self, params: Dict[str, B.Numeric], X: B.Numeric, X2: Optional[B.Numeric] = None, **kwargs # type: ignore 

227 ) -> B.Numeric: 

228 _check_field_in_params(params, "lengthscale") 

229 _check_1_vector(params["lengthscale"], 'params["lengthscale"]') 

230 

231 _check_field_in_params(params, "nu") 

232 _check_1_vector(params["nu"], 'params["nu"]') 

233 

234 weights = B.cast(B.dtype(params["nu"]), self.eigenvalues(params)) # [L, 1] 

235 Phi = self.eigenfunctions 

236 K = Phi.weighted_outerproduct(weights, X, X2, **kwargs) # [N, N2] 

237 if is_complex(K): 

238 return B.real(K) 

239 else: 

240 return K 

241 

242 def K_diag(self, params: Dict[str, B.Numeric], X: B.Numeric, **kwargs) -> B.Numeric: 

243 _check_field_in_params(params, "lengthscale") 

244 _check_1_vector(params["lengthscale"], 'params["lengthscale"]') 

245 

246 _check_field_in_params(params, "nu") 

247 _check_1_vector(params["nu"], 'params["nu"]') 

248 

249 weights = B.cast(B.dtype(params["nu"]), self.eigenvalues(params)) # [L, 1] 

250 Phi = self.eigenfunctions 

251 K_diag = Phi.weighted_outerproduct_diag(weights, X, **kwargs) # [N,] 

252 if is_complex(K_diag): 

253 return B.real(K_diag) 

254 else: 

255 return K_diag