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
« 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"""
7import abc
9import lab as B
10import numpy as np
11from beartype.typing import List, Optional, Tuple
13from geometric_kernels.spaces.base import DiscreteSpectrumSpace
14from geometric_kernels.spaces.eigenfunctions import EigenfunctionsWithAdditionTheorem
17class LieGroupCharacter(abc.ABC):
18 """
19 Class that represents a character of a Lie group.
20 """
22 @abc.abstractmethod
23 def __call__(self, gammas: B.Numeric) -> B.Numeric:
24 """
25 Compute the character on `gammas` lying in a maximal torus.
27 :param gammas:
28 [..., rank] where `rank` is the dimension of a maximal torus.
30 :return:
31 An array of shape [...] representing the values of the characters.
32 The values can be complex-valued.
33 """
34 raise NotImplementedError
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.
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.
59 Defaults to True.
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.
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.
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 """
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.
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.
106 :param num_levels:
107 Number of levels.
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
116 @abc.abstractmethod
117 def _compute_eigenvalue(self, signature: Tuple[int, ...]) -> B.Float:
118 """
119 Compute eigenvalue corresponding to `signature`.
121 :param signature:
122 The signature.
124 :return:
125 The eigenvalue.
126 """
127 raise NotImplementedError
129 @abc.abstractmethod
130 def _compute_dimension(self, signature: Tuple[int, ...]) -> int:
131 """
132 Compute dimension of the representation corresponding to `signature`.
134 :param signature:
135 The signature.
137 :return:
138 The dimension.
139 """
140 raise NotImplementedError
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`.
149 :param signature:
150 The signature.
152 :return:
153 The character, represented by :class:`LieGroupCharacter`.
154 """
155 raise NotImplementedError
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.
163 :param X:
164 An [N, n, n]-shaped array, a batch of N matrices of size nxn.
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
172 def inverse(self, X: B.Numeric) -> B.Numeric:
173 """
174 Should call the group's static
175 :meth:`~.spaces.CompactMatrixLieGroup.inverse` method.
177 :param X:
178 As in :meth:`~.spaces.CompactMatrixLieGroup.inverse`.
179 """
180 raise NotImplementedError
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`.
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.
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.
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)
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
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.
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.
227 Defaults to None, in which case X is used for X2.
228 :param ``**kwargs``:
229 Any additional parameters.
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]
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.
249 :param X:
250 As in :meth:`_addition_theorem`.
251 :param ``**kwargs``:
252 As in :meth:`_addition_theorem`.
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]
264 @property
265 def num_levels(self) -> int:
266 return self._num_levels
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
274 @property
275 def num_eigenfunctions_per_level(self) -> List[int]:
276 """
277 The number of eigenfunctions per level.
279 :return:
280 List of squares of dimensions of the unitary irreducible
281 representations.
282 """
283 return [d**2 for d in self._dimensions]
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.
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 """
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.
305 :param X:
306 A batch [..., n, n] of elements of the group.
307 Each element is a n x n matrix.
309 :return:
310 A batch [..., n, n] with each n x n matrix inverted.
311 """
312 raise NotImplementedError