Coverage for geometric_kernels/spaces/lie_groups.py: 87%

69 statements  

« prev     ^ index     » next       coverage.py v7.11.3, created at 2025-11-16 21:43 +0000

1""" 

2Abstract base interface for compact matrix Lie groups and its abstract 

3friends, a subclass of :class:`~.eigenfunctions.Eigenfunctions` and an 

4abstract base class for Lie group characters. 

5""" 

6 

7import abc 

8 

9import lab as B 

10import numpy as np 

11from beartype.typing import List, Optional, Tuple 

12 

13from geometric_kernels.spaces.base import DiscreteSpectrumSpace 

14from geometric_kernels.spaces.eigenfunctions import EigenfunctionsWithAdditionTheorem 

15 

16 

17class LieGroupCharacter(abc.ABC): 

18 """ 

19 Class that represents a character of a Lie group. 

20 """ 

21 

22 @abc.abstractmethod 

23 def __call__(self, gammas: B.Numeric) -> B.Numeric: 

24 """ 

25 Compute the character on `gammas` lying in a maximal torus. 

26 

27 :param gammas: 

28 [..., rank] where `rank` is the dimension of a maximal torus. 

29 

30 :return: 

31 An array of shape [...] representing the values of the characters. 

32 The values can be complex-valued. 

33 """ 

34 raise NotImplementedError 

35 

36 

37class WeylAdditionTheorem(EigenfunctionsWithAdditionTheorem): 

38 r""" 

39 A class for computing the sums of outer products of certain combinations 

40 (we call "levels") of Laplace-Beltrami eigenfunctions on compact Lie 

41 groups. These are much like the *zonal spherical harmonics* of the 

42 :class:`~.SphericalHarmonics` class. However, the key to computing them 

43 is representation-theoretic: they are proportional to *characters* of 

44 irreducible unitary representations of the group. These characters, in their 

45 turn, can be algebraically computed using the *Weyl character formula*. See 

46 :cite:t:`azangulov2024a` for the mathematical details behind this class. 

47 

48 :param n: 

49 The order of the Lie group, e.g. for SO(5) this is 5, for SU(3) this is 3. 

50 :param num_levels: 

51 The number of levels used for kernel approximation. Here, each level 

52 corresponds to an irreducible unitary representation of the group. 

53 :param compute_characters: 

54 Whether or not to actually compute the *characters*. Setting this parameter 

55 to False might make sense if you do not care about eigenfunctions (or sums 

56 of outer products thereof), but care about eigenvalues, dimensions of 

57 irreducible unitary representations, etc. 

58 

59 Defaults to True. 

60 

61 .. note:: 

62 Unlike :class:`~.SphericalHarmonics`, we do not expect the 

63 descendants of this class to compute the actual Laplace-Beltrami 

64 eigenfunctions (which are similar to (non-zonal) *spherical harmonics* 

65 of :class:`~.SphericalHarmonics`), implementing the `__call___` method. 

66 In this case, this is not really necessary and computationally harder. 

67 

68 .. note:: 

69 Unlike in :class:`~.SphericalHarmonics`, here the levels do not 

70 necessarily correspond to whole eigenspaces (all eigenfunctions that 

71 correspond to the same eigenvalue). Here, levels are defined in terms of 

72 the algebraic structure of the group: they are the eigenfunctions that 

73 correspond to the same irreducible unitary representation of the group. 

74 

75 .. note:: 

76 Here we break the general convention that the subclasses of the 

77 :class:`~.eigenfunctions.Eigenfunctions` only provide an interface for 

78 working with eigenfunctions, not eigenvalues, offering an interface 

79 for computing the latter as well. 

80 """ 

81 

82 def __init__(self, n: int, num_levels: int, compute_characters: bool = True): 

83 self._num_levels = num_levels 

84 self._signatures = self._generate_signatures(self._num_levels) 

85 self._eigenvalues = np.array( 

86 [self._compute_eigenvalue(signature) for signature in self._signatures] 

87 ) 

88 self._dimensions = [ 

89 self._compute_dimension(signature) for signature in self._signatures 

90 ] 

91 if compute_characters: 

92 self._characters = [ 

93 self._compute_character(n, signature) for signature in self._signatures 

94 ] 

95 self._num_eigenfunctions: Optional[int] = None # To be computed when needed. 

96 

97 @abc.abstractmethod 

98 def _generate_signatures(self, num_levels: int) -> List[Tuple[int, ...]]: 

99 """ 

100 Generate the signatures of `self.num_levels` irreducible unitary 

101 representations of the group that (likely) correspond to the smallest 

102 Laplace-Beltrami eigenvalues. These signatures index representations, 

103 mathematically they are vectors of dimension equal to the *rank* of the 

104 group, which in its turn is the dimension of maximal tori. 

105 

106 :param num_levels: 

107 Number of levels. 

108 

109 :return: 

110 List of signatures of irreducible unitary representations. Each 

111 signature is itself a tuple of integers, whose length is equal to 

112 the rank of the group typically stored in `self.rank`. 

113 """ 

114 raise NotImplementedError 

115 

116 @abc.abstractmethod 

117 def _compute_eigenvalue(self, signature: Tuple[int, ...]) -> B.Float: 

118 """ 

119 Compute eigenvalue corresponding to `signature`. 

120 

121 :param signature: 

122 The signature. 

123 

124 :return: 

125 The eigenvalue. 

126 """ 

127 raise NotImplementedError 

128 

129 @abc.abstractmethod 

130 def _compute_dimension(self, signature: Tuple[int, ...]) -> int: 

131 """ 

132 Compute dimension of the representation corresponding to `signature`. 

133 

134 :param signature: 

135 The signature. 

136 

137 :return: 

138 The dimension. 

139 """ 

140 raise NotImplementedError 

141 

142 @abc.abstractmethod 

143 def _compute_character( 

144 self, n: int, signature: Tuple[int, ...] 

145 ) -> LieGroupCharacter: 

146 """ 

147 Compute character of the representation corresponding to `signature`. 

148 

149 :param signature: 

150 The signature. 

151 

152 :return: 

153 The character, represented by :class:`LieGroupCharacter`. 

154 """ 

155 raise NotImplementedError 

156 

157 @abc.abstractmethod 

158 def _torus_representative(self, X: B.Numeric) -> B.Numeric: 

159 """ 

160 The function maps a Lie group element `X` to a maximal torus of the 

161 Lie group. 

162 

163 :param X: 

164 An [N, n, n]-shaped array, a batch of N matrices of size nxn. 

165 

166 :return: 

167 An [N, R]-shape array, a batch of N vectors of size R, where R is 

168 the dimension of maximal tori, i.e. the rank of the Lie group. 

169 """ 

170 raise NotImplementedError 

171 

172 def inverse(self, X: B.Numeric) -> B.Numeric: 

173 """ 

174 Should call the group's static 

175 :meth:`~.spaces.CompactMatrixLieGroup.inverse` method. 

176 

177 :param X: 

178 As in :meth:`~.spaces.CompactMatrixLieGroup.inverse`. 

179 """ 

180 raise NotImplementedError 

181 

182 def _difference(self, X: B.Numeric, X2: B.Numeric) -> B.Numeric: 

183 r""" 

184 Pairwise difference (in the group sense) between elements of the 

185 two batches, `X` and `X2`. 

186 

187 :param X: 

188 An [N, n, n]-shaped array, a batch of N matrices of size nxn. 

189 :param X2: 

190 An [N2, n, n]-shaped array, a batch of N2 matrices of size nxn. 

191 

192 :return: 

193 An [N, N2, n, n]-shaped array, all pairwise "differences": 

194 X1[j, :, :] * inv(X2[i, :, :]) for all 0 <= i < N, 0 <= j < N2. 

195 

196 .. note:: 

197 Doing X1[j, :, :] * inv(X2[i, :, :]) is as permissible as 

198 doing inv(X2[i, :, :]) * X1[j, :, :] which is actually used in 

199 :cite:t:`azangulov2024a`. This is because $\chi(x y x^{-1})=\chi(y)$ 

200 which implies that $\chi(x y) = \chi(y x)$. 

201 """ 

202 X2_inv = self.inverse(X2) 

203 X_ = B.tile(X[..., None, :, :], 1, X2_inv.shape[0], 1, 1) # (N, N2, n, n) 

204 X2_inv_ = B.tile(X2_inv[None, ..., :, :], X.shape[0], 1, 1, 1) # (N, N2, n, n) 

205 

206 diff = B.reshape( 

207 B.matmul(X_, X2_inv_), X.shape[0], X2_inv.shape[0], X.shape[-1], X.shape[-1] 

208 ) # (N, N2, n, n) 

209 return diff 

210 

211 def _addition_theorem( 

212 self, X: B.Numeric, X2: Optional[B.Numeric] = None, **kwargs 

213 ) -> B.Numeric: 

214 r""" 

215 For each level (that corresponds to a unitary irreducible 

216 representation of the group), computes the sum of outer products of 

217 Laplace-Beltrami eigenfunctions that correspond to this level 

218 (representation). Uses the fact that such sums are equal to the 

219 character of the representation multiplied by the dimension of that 

220 representation. See :cite:t:`azangulov2024a` for mathematical details. 

221 

222 :param X: 

223 An [N, n, n]-shaped array, a batch of N matrices of size nxn. 

224 :param X2: 

225 An [N2, n, n]-shaped array, a batch of N2 matrices of size nxn. 

226 

227 Defaults to None, in which case X is used for X2. 

228 :param ``**kwargs``: 

229 Any additional parameters. 

230 

231 :return: 

232 An array of shape [N, N2, L]. 

233 """ 

234 if X2 is None: 

235 X2 = X 

236 diff = self._difference(X, X2) 

237 torus_repr_diff = self._torus_representative(diff) 

238 values = [ 

239 repr_dim * chi(torus_repr_diff)[..., None] # [N, N2, 1] 

240 for chi, repr_dim in zip(self._characters, self._dimensions) 

241 ] 

242 return B.concat(*values, axis=-1) # [N, N2, L] 

243 

244 def _addition_theorem_diag(self, X: B.Numeric, **kwargs) -> B.Numeric: 

245 """ 

246 A more efficient way of computing the diagonals of the matrices 

247 `self._addition_theorem(X, X)[:, :, l]` for all l from 0 to L-1. 

248 

249 :param X: 

250 As in :meth:`_addition_theorem`. 

251 :param ``**kwargs``: 

252 As in :meth:`_addition_theorem`. 

253 

254 :return: 

255 An array of shape [N, L]. 

256 """ 

257 ones = B.ones(B.dtype(X), *X.shape[:-2], 1) 

258 values = [ 

259 repr_dim * repr_dim * ones # [N, 1], because chi(X*inv(X))=repr_dim 

260 for repr_dim in self._dimensions 

261 ] 

262 return B.concat(*values, axis=1) # [N, L] 

263 

264 @property 

265 def num_levels(self) -> int: 

266 return self._num_levels 

267 

268 @property 

269 def num_eigenfunctions(self) -> int: 

270 if self._num_eigenfunctions is None: 

271 self._num_eigenfunctions = sum(self.num_eigenfunctions_per_level) 

272 return self._num_eigenfunctions 

273 

274 @property 

275 def num_eigenfunctions_per_level(self) -> List[int]: 

276 """ 

277 The number of eigenfunctions per level. 

278 

279 :return: 

280 List of squares of dimensions of the unitary irreducible 

281 representations. 

282 """ 

283 return [d**2 for d in self._dimensions] 

284 

285 

286class CompactMatrixLieGroup(DiscreteSpectrumSpace): 

287 r""" 

288 A base class for compact Lie groups of matrices, subgroups of the 

289 real or complex general linear group GL(n), n being some integer. 

290 

291 The group operation is the standard matrix multiplication, and the group 

292 inverse is the standard matrix inverse. Despite this, we make the 

293 subclasses implement their own inverse routine, because in special cases 

294 it can typically be implemented much more efficiently. For example, for 

295 the special orthogonal group :class:`~.spaces.SpecialOrthogonal`, the 

296 inverse is equivalent to a simple transposition. 

297 """ 

298 

299 @staticmethod 

300 @abc.abstractmethod 

301 def inverse(X: B.Numeric) -> B.Numeric: 

302 """ 

303 Inverse of a batch `X` of the elements of the group. A static method. 

304 

305 :param X: 

306 A batch [..., n, n] of elements of the group. 

307 Each element is a n x n matrix. 

308 

309 :return: 

310 A batch [..., n, n] with each n x n matrix inverted. 

311 """ 

312 raise NotImplementedError