Coverage for geometric_kernels/spaces/graph_edges.py: 70%
247 statements
« prev ^ index » next coverage.py v7.11.3, created at 2025-11-16 21:43 +0000
1"""
2This module provides the :class:`EdgeGraph` space.
3"""
5import lab as B
6import numpy as np
7from beartype.typing import Dict, List, Optional, Tuple, Union
8from scipy.sparse import csr_matrix, issparse, lil_matrix
10from geometric_kernels.lab_extras import (
11 SparseArray,
12 count_nonzero,
13 dtype_integer,
14 eigenpairs,
15 float_like,
16 int_like,
17 take_along_axis,
18)
19from geometric_kernels.spaces.base import HodgeDiscreteSpectrumSpace
20from geometric_kernels.spaces.eigenfunctions import (
21 Eigenfunctions,
22 EigenfunctionsFromEigenvectors,
23)
24from geometric_kernels.utils.utils import _check_matrix, _check_rank_1_array
27class GraphEdges(HodgeDiscreteSpectrumSpace):
28 """
29 The GeometricKernels space representing the edge set of a user-provided
30 graph. The graph must be unweighted, undirected, and without loops.
31 However, we assume the graph is oriented: for every edge, either (i, j)
32 or (j, i) is chosen as the positively oriented edge, with the other one
33 being negatively oriented. Any function f on this space must satisfy
34 f(e) = -f(-e) if e is a positively oriented edge and -e is the
35 corresponding negatively oriented edge.
37 All positively oriented edges are indexed **from 1 to the number of edges
38 (inclusive)**. The reason why we start indexing the edges from 1 is to
39 allow for edge -e to correspond to the opposite orientation of edge e.
40 Importantly, orientation or a particular ordering of edges does not
41 affect the geometry of the space. However, changing those requires the
42 respective change in the way you represent the data.
44 For this space, each individual eigenfunction constitutes a *level*. If
45 possible, we recommend using num=num_edges levels. Since the current
46 implementation does not support sparse eigensolvers anyway, this should
47 not create too much time/memory overhead during the precomputations and
48 will also ensure that all types of eigenfunctions are present in the
49 eigensystem.
51 .. admonition:: Tutorial
53 A tutorial on how to use this space is available in the
54 :doc:`GraphEdges.ipynb </examples/GraphEdges>` notebook.
56 .. admonition:: Theory
58 Mathematically, this space is actually a simplicial 2-complex, and
59 the functions satisfying the condition f(e) = -f(-e) are called
60 1-cochains. These admit a Hodge decomposition into the harmonic,
61 gradient, and curl parts. You can find more about Hodge decomposition
62 on :doc:`this page </theory/graphs>`.
64 .. admonition:: Construction
66 **Easy way:** construct the space from the adjacency matrix of the graph.
68 The easiest way to construct an instance of this space is to use the
69 :meth:`from_adjacency` method, which constructs the space from the
70 adjacency matrix of the graph. By default, this would use all possible
71 triangles in the graph, but the user can also provide a list of
72 triangles explicitly. These triangles define the curl operator on the
73 graph and are thus instrumental in defining the Hodge decomposition.
74 If you don't know what triangles to provide, you can leave this
75 parameter empty, and the space will compute the maximal set of
76 triangles for you, which is a good solution in most cases.
78 When constructing an instance this way, the orientation of the edges and
79 triangles is chosen automatically in such a way that (i, j) is always
80 positively oriented if i < j, and the triangles are of form (e1, e2, e3)
81 where e1 = (i, j), e2 = (j, k), e3 = -(i, k), and i < j < k.
83 **Hard way:** construct the space from the oriented edges and triangles.
85 When constructing an instance of this space directly via :meth:`__init__`,
86 the user provides an `oriented_edges` array, mapping edge indices to
87 ordered pairs of node indices, with order defining the positive
88 orientation, and an `oriented_triangles` array, mapping triangle indices
89 to triples of signed edge indices. These arrays must be compatible in
90 the sense that the edges in triangles must be connected. This way, the
91 user has full control over the orientation and ordering of the edges.
93 .. admonition:: Complexity
95 The current implementation of the GraphEdges space is supposed to occupy
96 O(num_edges^2) memory and take O(num_edges^3) time to compute the eigensystem.
98 Currently, it does not support sparse eigensolvers, which would allow
99 storing the Laplacian matrices in a sparse format and potentially reduce
100 the O(num_edges^3) complexity to O(num_edges) with a large constant.
102 :param num_nodes:
103 The number of nodes in the graph.
105 :param oriented_edges:
106 A 2D array of shape [num_edges, 2], where num_edges is the number of edges
107 in the graph. Each row of the array represents an oriented edge, i.e.,
108 a pair of node indices. If `oriented_edges[e]` is `[i, j]`, then the
109 edge with index e+1 (edge indexing starts from 1) represents an edge
110 from node `i` to node `j` which is considered positively oriented, while
111 the edge with index -e-1 represents an edge from node `j` to node `i`.
113 Make sure that `oriented_edges` and `oriented_triangles` are of the
114 backend (NumPy / JAX / TensorFlow, PyTorch) that you wish to use for
115 internal computations.
117 :param oriented_triangles:
118 A 2D array of shape [num_triangles, 3], where num_triangles is the number
119 of triangles in the graph. Each row of the array represents an
120 oriented triangle, i.e., a triple of signed edge indices. If
121 `oriented_triangles[t]` is `[i, j, k]`, then `t` consists of the edges
122 `|i|`, `|j|`, and `|k|`, where the sign of the edge index indicates
123 the orientation of the edge. Optional. If not provided or set to
124 `None`, we assume that the set of triangles is empty.
126 Make sure that `oriented_edges` and `oriented_triangles` are of the
127 backend (NumPy / JAX / TensorFlow, PyTorch) that you wish to use for
128 internal computations.
130 :param checks_mode:
131 A keyword argument that determines the level of checks performed on
132 the oriented_edges and oriented_triangles arrays. The default value is
133 "simple", which performs only basic checks. The other possible values
134 are "comprehensive" and "skip", which perform more extensive checks
135 and no checks, respectively. The "comprehensive" mode is useful for
136 debugging but can be slow for large graphs.
138 :param index:
139 Optional. A sparse matrix of shape [num_nodes, num_nodes] such that
140 `index[i, j]` is the index of the edge `(i, j)` in the `oriented_edges`.
141 `index[i, j]` is positive if (i, j) corresponds to positive orientation
142 of the edge, and negative if it corresponds to the negative orientation.
143 If provided, should be compatible with the `oriented_edges` array.
145 .. admonition:: Citation
147 If you use this GeometricKernels space in your research, please consider
148 citing :cite:t:`yang2024`.
149 """
151 def __init__(
152 self,
153 num_nodes: int,
154 oriented_edges: B.Int,
155 oriented_triangles: Optional[B.Int],
156 *,
157 checks_mode: str = "simple",
158 index: Optional[csr_matrix] = None,
159 ):
160 if checks_mode not in ["simple", "comprehensive", "skip"]:
161 raise ValueError(
162 "The `checks_mode` parameter must be 'simple', 'comprehensive', or 'skip'."
163 )
164 else:
165 do_checks = checks_mode != "skip"
166 comprehensive_checks = checks_mode == "comprehensive"
168 self.cache: Dict[int, Tuple[B.Numeric, B.Numeric]] = {}
169 self.num_nodes = num_nodes
171 if do_checks:
172 self._checks_oriented_edges(oriented_edges, num_nodes, comprehensive_checks)
173 self.num_edges = oriented_edges.shape[0]
174 self.oriented_edges = oriented_edges
176 if oriented_triangles is not None:
177 if do_checks:
178 self._checks_oriented_triangles(
179 oriented_triangles, comprehensive_checks
180 )
181 if comprehensive_checks:
182 self._checks_compatible(oriented_triangles)
183 self.num_triangles = oriented_triangles.shape[0]
184 else:
185 self.num_triangles = 0
186 self.oriented_triangles = oriented_triangles
188 if index is not None:
189 if comprehensive_checks:
190 self._check_index(index)
191 self.index = index
192 else:
193 self.index = self._compute_index(num_nodes, oriented_edges)
195 self._set_laplacian()
197 def __str__(self):
198 return f"GraphEdges({self.num_edges})"
200 @staticmethod
201 def _compute_index(num_nodes: int, oriented_edges: B.Numeric) -> csr_matrix:
202 """
203 Construct the index matrix from the oriented edges.
205 :param num_nodes:
206 The number of nodes in the graph.
208 :param oriented_edges:
209 The oriented edges array.
211 :return:
212 The index matrix. A scipy csr_matrix of shape [num_nodes, num_nodes] such
213 that `index[i, j]` is the index of the edge `(i, j)`.
214 """
215 result = lil_matrix((num_nodes, num_nodes), dtype=int)
216 for i in range(oriented_edges.shape[0]):
217 result[oriented_edges[i, 0], oriented_edges[i, 1]] = i + 1
218 result[oriented_edges[i, 1], oriented_edges[i, 0]] = -i - 1
219 return result.tocsr()
221 def _checks_oriented_edges( # NOQA: C901
222 self, oriented_edges: B.Numeric, num_nodes: int, comprehensive: bool = False
223 ):
224 """
225 Checks if `oriented_edges` is of appropriate structure.
227 :param oriented_edges:
228 The oriented edges array.
230 :param comprehensive:
231 If True, perform more extensive checks.
232 """
233 _check_matrix(oriented_edges, "oriented_edges")
235 if B.shape(oriented_edges)[1] != 2:
236 raise ValueError("`oriented_edges` must have shape (*, 2).")
238 if B.dtype(oriented_edges) != int_like(oriented_edges):
239 raise ValueError("`oriented_edges` must be an array of integers.")
240 if B.any(oriented_edges < 0):
241 raise ValueError("`oriented_edges` must contain only non-negative values.")
242 if B.any(oriented_edges >= self.num_nodes):
243 raise ValueError(
244 "The values in the `oriented_edges` array must be less than `self.num_nodes.`"
245 )
246 if B.any(oriented_edges[:, 0] - oriented_edges[:, 1] == 0):
247 raise ValueError("Loops are not allowed.")
249 if not comprehensive:
250 return
252 num_edges = oriented_edges.shape[0]
254 for i in range(num_edges):
255 for j in range(i + 1, num_edges):
256 if B.all(oriented_edges[i, :] == oriented_edges[j, :]):
257 raise ValueError(
258 "`oriented_edges` must not contain duplicate edges."
259 )
260 if B.all(oriented_edges[i, :] == oriented_edges[j, ::-1]):
261 raise ValueError(
262 "`oriented_edges` must not contain duplicate edges."
263 )
265 if set(range(self.num_nodes)) != set(B.to_numpy(B.flatten(oriented_edges))):
266 raise ValueError("`oriented_edges` must contain all nodes.")
268 def _checks_oriented_triangles( # NOQA: C901
269 self, oriented_triangles: B.Numeric, comprehensive=False
270 ):
271 """
272 Checks if `oriented_triangles` is of appropriate structure.
274 :param oriented_triangles:
275 The oriented triangles array.
277 :param comprehensive:
278 If True, perform more extensive checks.
279 """
280 _check_matrix(oriented_triangles, "oriented_triangles")
281 if B.shape(oriented_triangles)[1] != 3:
282 raise ValueError("`oriented_triangles` must have shape (*, 3).")
284 if B.dtype(oriented_triangles) != int_like(oriented_triangles):
285 raise ValueError("`oriented_triangles` must be an array of integers.")
286 if B.any(B.abs(oriented_triangles) < 1):
287 raise ValueError("`oriented_triangles` must contain only non-zero values.")
288 if B.any(B.abs(oriented_triangles) > self.num_edges):
289 raise ValueError(
290 "Absolute values in `oriented_triangles` array must be less than or equal to `self.num_edges`."
291 )
293 if (
294 B.any(
295 B.abs(oriented_triangles[:, 0]) - B.abs(oriented_triangles[:, 1]) == 0
296 )
297 and B.any(
298 B.abs(oriented_triangles[:, 0]) - B.abs(oriented_triangles[:, 2]) == 0
299 )
300 and B.any(
301 B.abs(oriented_triangles[:, 1]) - B.abs(oriented_triangles[:, 2]) == 0
302 )
303 ):
304 raise ValueError("Triangles must consist of 3 different edges.")
306 if not comprehensive:
307 return
309 num_triangles = oriented_triangles.shape[0]
311 for i in range(num_triangles):
312 for j in range(i + 1, num_triangles):
313 if B.all(oriented_triangles[i, :] == oriented_triangles[j, :]):
314 raise ValueError(
315 "The oriented_triangles array must not contain duplicate triangles."
316 )
318 def _checks_compatible(
319 self,
320 oriented_triangles: B.Numeric,
321 ):
322 """
323 Checks if `self.oriented_edges` and `oriented_triangles` are compatible.
325 :param oriented_triangles:
326 The oriented triangles array.
327 """
329 if B.dtype(self.oriented_edges) != B.dtype(oriented_triangles):
330 raise ValueError(
331 "`oriented_edges` and `oriented_triangles` must have the same dtype."
332 )
334 num_triangles = oriented_triangles.shape[0]
335 for t in range(num_triangles):
336 resolved_edges = self.resolve_edges(oriented_triangles[t, :])
337 if resolved_edges[0, 1] != resolved_edges[1, 0]:
338 raise ValueError("The edges in the triangle must be connected.")
339 if resolved_edges[1, 1] != resolved_edges[2, 0]:
340 raise ValueError("The edges in the triangle must be connected.")
341 if resolved_edges[2, 1] != resolved_edges[0, 0]:
342 raise ValueError("The edges in the triangle must be connected.")
344 def _check_index(self, index: csr_matrix):
345 edges = []
346 for e in range(1, self.oriented_edges.shape[0] + 1):
347 i, j = self.oriented_edges[e - 1, :]
348 if index[i, j] != e:
349 raise ValueError("`index` must be compatible with `oriented_edges`.")
350 if index[j, i] != -e:
351 raise ValueError("`index` must be compatible with `oriented_edges`.")
352 edges.append((min(i, j), max(i, j)))
354 for i in range(self.num_nodes):
355 for j in range(i + 1, self.num_nodes):
356 if (i, j) not in edges:
357 if index[i, j] != 0:
358 raise ValueError(
359 "`index` must be compatible with `oriented_edges`."
360 )
361 if index[j, i] != 0:
362 raise ValueError(
363 "`index` must be compatible with `oriented_edges`."
364 )
366 def resolve_edges(self, es: B.Int) -> B.Int:
367 r"""
368 Resolve the signed edge indices to node indices.
370 :param es:
371 A 1-dimensional array of edge indices. Each edge index is a number
372 from 1 to the number of edges (incl.).
374 :return:
375 A 2-dimensional array `result` such that `result[e, :]` is `[i, j]`
376 where \|e\| = (i, j) if e > 0 and \|e\| = (j, i) if e < 0.
377 """
378 _check_rank_1_array(es, "es")
379 if not (B.all(B.abs(es) >= 1) and B.all(B.abs(es) <= self.num_edges)):
380 raise ValueError("`abs(es)` must lie in the interval [1, `num_edges`].")
382 result = self.oriented_edges[B.abs(es) - 1]
383 result = B.where(B.expand_dims(es > 0, axis=-1), result, result[:, ::-1])
384 return result
386 def resolve_triangles(self, ts: B.Int) -> B.Int:
387 """
388 Resolve the triangle indices to node indices.
390 :param ts:
391 A 1-dimensional array of triangle indices. Each triangle index is a
392 number from 0 to the number of triangles.
394 :return:
395 A 3-dimensional array `result` such that `result[t, :]` is `[i, j, k]`
396 where i = e1[0], j = e2[0], k = e3[0], and e1, e2, e3 are the
397 oriented edges constituting the triangle `t`.
398 """
399 _check_rank_1_array(ts, "ts")
400 if not (B.all(B.abs(ts) >= 0) and B.all(B.abs(ts) < self.num_triangles)):
401 raise ValueError("`abs(ts)` must lie in the interval [1, `num_edges`].")
403 edge_indices = B.flatten(
404 self.oriented_triangles[ts]
405 ) # [N,] -> [N, 3] -> [N*3,]
406 edges = B.reshape(
407 self.resolve_edges(edge_indices), len(ts), 3, 2
408 ) # [N*3,] -> [N*3, 2] -> [N, 3, 2]
409 return edges[:, :, 0] # [N, 3, 2] -> [N, 3]
411 @classmethod
412 def from_adjacency( # noqa: C901
413 cls,
414 adjacency_matrix: Union[B.NPNumeric, SparseArray],
415 type_reference: B.RandomState,
416 *,
417 triangles: Optional[List[Tuple[int, int, int]]] = None,
418 checks_mode: str = "simple",
419 ) -> "GraphEdges":
420 """
421 Construct the GraphEdges space from the adjacency matrix of a graph.
423 :param adjacency_matrix:
424 Adjacency matrix of a graph. A numpy array of shape
425 `[num_nodes, num_nodes]` where `num_nodes` is the number of nodes in the
426 graph. `adjacency_matrix[i, j]` can only be 0 or 1.
428 :param type_reference:
429 A random state object of the preferred backend to infer backend from.
431 :param triangles:
432 A list of tuples of three integers representing the nodes of the
433 triangles in the graph. If not provided or set to None, the maximal
434 possible set of triangles will be computed and used.
436 :param checks_mode:
437 Forwards the `checks_mode` parameter to the constructor.
439 :return:
440 A constructed instance of the GraphEdges space.
441 """
442 if isinstance(adjacency_matrix, np.ndarray):
443 index = csr_matrix(adjacency_matrix, dtype=int)
444 elif issparse(adjacency_matrix):
445 index = csr_matrix(adjacency_matrix, dtype=int, copy=True)
446 else:
447 raise ValueError(
448 f"The adjacency matrix must be a numpy array or a scipy sparse matrix not {type(adjacency_matrix)}. Use `type_reference` to specify the backend."
449 )
451 _check_matrix(index, "adjacency_matrix")
452 if B.shape(index)[0] != B.shape(index)[1]:
453 raise ValueError("`adjacency_matrix` must be a square matrix.")
454 if (abs(index - index.T) > 1e-10).nnz != 0:
455 raise ValueError("`adjacency_matrix` must be symmetric.")
456 if (index.diagonal() != 0).any():
457 raise ValueError("`adjacency_matrix` must have zeros on the diagonal.")
458 if np.sum(index.data == 1) + np.sum(index.data == 0) != len(index.data):
459 raise ValueError("`adjacency_matrix` can only contain zeros and ones.")
461 number_of_nodes = index.shape[0]
462 number_of_edges = np.sum(index.data) // 2
464 oriented_edges = np.zeros((number_of_edges, 2), dtype=np.int32)
466 cur_edge_ind = 1
467 for i in range(number_of_nodes):
468 for j in range(i + 1, number_of_nodes):
469 if index[i, j] == 1:
470 oriented_edges[cur_edge_ind - 1, 0] = i
471 oriented_edges[cur_edge_ind - 1, 1] = j
472 # We also store cur_edge_ind in the index matrix for later use
473 index[i, j] = cur_edge_ind
474 index[j, i] = -index[i, j]
475 cur_edge_ind += 1
476 if cur_edge_ind != number_of_edges + 1:
477 # double check that we have the right number of edges
478 raise RuntimeError(
479 "This should have never happened, please report a bug at https://github.com/geometric-kernels/GeometricKernels/issues."
480 )
481 oriented_edges = B.cast(dtype_integer(type_reference), oriented_edges)
483 if triangles is None:
484 triangles = []
485 for i in range(number_of_nodes):
486 for j in range(i + 1, number_of_nodes):
487 for k in range(j + 1, number_of_nodes):
488 if index[i, j] != 0 and index[j, k] != 0 and index[i, k] != 0:
489 triangles.append((i, j, k))
491 number_of_triangles = len(triangles)
492 oriented_triangles = np.zeros((number_of_triangles, 3), dtype=np.int32)
493 for triangle_index, triangle in enumerate(triangles):
494 i, j, k = sorted(triangle) # sort the nodes in the triangle
495 oriented_triangles[triangle_index, 0] = index[i, j]
496 oriented_triangles[triangle_index, 1] = index[j, k]
497 oriented_triangles[triangle_index, 2] = index[k, i]
498 oriented_triangles = B.cast(dtype_integer(type_reference), oriented_triangles)
500 return cls(
501 number_of_nodes,
502 oriented_edges,
503 oriented_triangles,
504 checks_mode=checks_mode,
505 index=index,
506 )
508 @property
509 def dimension(self) -> int:
510 """
511 :return:
512 0.
513 """
514 return 0 # this is needed for the kernel math to work out
516 def _set_laplacian(self):
517 """
518 Construct the appropriate graph Laplacian from the adjacency matrix.
519 """
521 # This does node_to_edge_incidence.T @ node_to_edge_incidence
522 # TODO: make this more efficient, avoid for loops.
523 self._down_laplacian = np.zeros(
524 (self.num_edges, self.num_edges), dtype=np.int32
525 )
526 for i in range(self.num_edges):
527 for j in range(self.num_edges):
528 self._down_laplacian[i, j] = (
529 int(self.oriented_edges[i, 0] == self.oriented_edges[j, 0])
530 + int(self.oriented_edges[i, 1] == self.oriented_edges[j, 1])
531 - int(self.oriented_edges[i, 0] == self.oriented_edges[j, 1])
532 - int(self.oriented_edges[i, 1] == self.oriented_edges[j, 0])
533 )
534 self._down_laplacian = B.cast(
535 float_like(self.oriented_edges), self._down_laplacian
536 )
538 # This does edge_to_triangle_incidence @ edge_to_triangle_incidence.T
539 # TODO: make this more efficient, avoid for loops.
540 self._up_laplacian = np.zeros((self.num_edges, self.num_edges), dtype=np.int32)
541 for i in range(self.num_edges):
542 for j in range(self.num_edges):
543 for t in range(self.num_triangles):
544 cur_triangle = set(B.to_numpy(self.oriented_triangles[t, :]))
545 if (i + 1 in cur_triangle and j + 1 in cur_triangle) or (
546 -i - 1 in cur_triangle and -j - 1 in cur_triangle
547 ):
548 self._up_laplacian[i, j] += 1
549 if (i + 1 in cur_triangle and -j - 1 in cur_triangle) or (
550 -i - 1 in cur_triangle and j + 1 in cur_triangle
551 ):
552 self._up_laplacian[i, j] += -1
553 self._up_laplacian = B.cast(float_like(self.oriented_edges), self._up_laplacian)
555 self._hodge_laplacian = self._down_laplacian + self._up_laplacian
557 def _get_eigensystem(self, num):
558 """
559 Returns the first `num` eigenvalues and eigenvectors of the Hodge Laplacian.
560 Caches the solution to prevent re-computing the same values.
562 :param num:
563 Number of eigenpairs to return. Performs the computation at the
564 first call. Afterwards, fetches the result from cache.
566 :return:
567 A tuple of eigenvectors [n, num], eigenvalues [num, 1].
568 """
570 # The current implementation computes the complete set of eigenpairs of
571 # three num_edges x num_edges Laplacian matrices. Since we don't currently
572 # support sparse Laplacians, this should not be a major bottleneck.
573 #
574 # Here are some important points for anyone who wants to add proper
575 # support for sparse Laplacians and sparse eigensolvers which allow
576 # computing only a few of the most important eigenpairs, like
577 # scipy.sparse.linalg.eigsh.
578 #
579 # 1. You cannot just compute the eigenpairs of the Hodge Laplacian and
580 # then filter them into the harmonic, gradient, or curl parts. This
581 # is because up and down edge Laplacians may have common non-zero
582 # eigenvalues. This means that the eigenspaces of the Hodge
583 # Laplacian that correspond to such eigenvalues will contain both
584 # gradient and curl types of eigenvectors. When an eigenspace is
585 # multi-dimensional, an eigensolver can return an arbitrary basis
586 # for it. Thus, for example, you may get a basis where all the vectors
587 # have non-zero grad and curl components at the same time. Thus, you
588 # won't be able to filter them.
589 #
590 # If you put in more effort, you can take the basis for each non-zero
591 # eigenvalue with multiplicity > 1, project each vector onto the
592 # pure-gradient and pure-curl subspaces (using up and down Laplacians),
593 # and then orthogonalize the projected vectors.
594 #
595 # 2. You cannot just request `num` eigenpairs from the Hodge Laplacian,
596 # up Laplacian and down Laplacian. Imagine that you have a space
597 # with 5 harmonic, 100 gradient, and 100 curl eigenpairs. Then,
598 # the first 105 eigenpairs of the up Laplacian will correspond to
599 # harmonic and gradient eigenvectors, both corresponding to zero
600 # eigenvalues, while only the last 100 eigenpairs will correspond
601 # to the actual curl eigenvectors that we are interested in getting
602 # from the up Laplacian. The same is true for the down Laplacian.
603 #
604 # 3. An elegant solution would be to construct
605 # > diff_laplacian = up_laplacian - down_laplacian
606 # and run a sparse eigensolver on this matrix, requesting `num`
607 # eigenpairs with the lowest absolute value of eigenvalues. Then,
608 # you get eigenpairs of the Hodge Laplacian corresponding to the
609 # first `num` smallest eigenvalues, but the gradient and curl
610 # eigenvectors would correspond to eigenvalues of different signs
611 # and thus would be separated.
613 eps = 1e-6
615 if num not in self.cache:
616 # We use Hodge Laplacian to find the eigenpairs of harmonic type,
617 # these are the ones associated to zero eigenvalues.
618 evals_hodge, evecs_hodge = eigenpairs(self._hodge_laplacian, self.num_edges)
620 # We use up and down Laplacians to find the eigenpairs of curl and
621 # gradient types. These are the ones associated to non-zero eigenvalues.
622 evals_up, evecs_up = eigenpairs(self._up_laplacian, self.num_edges)
623 evals_down, evecs_down = eigenpairs(self._down_laplacian, self.num_edges)
625 # We count the number of eigenpairs of each type for future reference.
626 n_harm = count_nonzero(evals_hodge < eps)
627 n_curl, n_grad = count_nonzero(evals_up >= eps), count_nonzero(
628 evals_down >= eps
629 )
631 # We concatenate the eigenvalues and eigenvectors of harmonic, curl, and
632 # gradient types into a single array.
633 evals = B.concat(
634 evals_hodge[evals_hodge < eps],
635 evals_up[evals_up >= eps],
636 evals_down[evals_down >= eps],
637 )
639 evecs = B.concat(
640 B.reshape(
641 evecs_hodge[
642 B.broadcast_to(
643 evals_hodge < eps, self.num_edges, self.num_edges
644 )
645 ],
646 self.num_edges,
647 -1,
648 ),
649 B.reshape(
650 evecs_up[
651 B.broadcast_to(evals_up >= eps, self.num_edges, self.num_edges)
652 ],
653 self.num_edges,
654 -1,
655 ),
656 B.reshape(
657 evecs_down[
658 B.broadcast_to(
659 evals_down >= eps, self.num_edges, self.num_edges
660 )
661 ],
662 self.num_edges,
663 -1,
664 ),
665 axis=-1,
666 )
668 evecs *= B.sqrt(self.num_edges) # to make sure average variance is 1.
670 # We get the indices that would sort the eigenvalues in ascending order.
671 sorted_indices = B.argsort(evals)
672 # And the indices that would inverse the sorting using sorted_indices.
673 sorted_sorted_indices = B.argsort(sorted_indices)
675 # In harmonic_idx, gradient_idx, and curl_idx, we store the indices of
676 # the eigenvalues and eigenvectors of harmonic, gradient, and curl types.
677 # We also make sure that the indices are not greater than `num`.
678 harmonic_idx = sorted_sorted_indices[:n_harm]
679 harmonic_idx = harmonic_idx[harmonic_idx < num]
680 curl_idx = sorted_sorted_indices[n_harm : n_harm + n_curl]
681 curl_idx = curl_idx[curl_idx < num]
682 grad_idx = sorted_sorted_indices[n_harm + n_curl : n_harm + n_curl + n_grad]
683 grad_idx = grad_idx[grad_idx < num]
685 # We sort the eigenvalues and eigenvectors in ascending order of
686 # eigenvalues, keeping only `num` of the first eigenpairs.
687 sorted_indices = sorted_indices[:num]
689 evals = take_along_axis(evals, sorted_indices)
690 evecs = take_along_axis(evecs, sorted_indices[None, :], axis=-1)
692 self.cache[num] = {
693 "evals": evals,
694 "evecs": evecs,
695 "harmonic_idx": harmonic_idx,
696 "gradient_idx": curl_idx,
697 "curl_idx": grad_idx,
698 }
700 return self.cache[num]
702 def get_number_of_eigenpairs(self, num: int, hodge_type: str) -> int:
703 """
704 Returns the number of eigenpairs of a particular type.
706 :param num:
707 Number of levels.
709 :param hodge_type:
710 The type of the eigenpairs. It can be "harmonic", "gradient", or "curl".
712 :return:
713 The number of eigenpairs of the specified type.
714 """
715 eigensystem = self._get_eigensystem(num)
716 return len(eigensystem[f"{hodge_type}_idx"])
718 def get_eigenfunctions(
719 self, num: int, hodge_type: Optional[str] = None
720 ) -> Eigenfunctions:
721 """
722 Returns the :class:`~.EigenfunctionsFromEigenvectors` object with
723 `num` levels (i.e., in this case, `num` eigenpairs) of a particular type.
725 :param num:
726 Number of levels.
728 :param hodge_type:
729 The type of the eigenbasis. It can be "harmonic", "gradient", or "curl".
731 :return:
732 EigenfunctionsFromEigenvectors object.
734 .. note::
735 The notion of *levels* is discussed in the documentation of the
736 :class:`~.kernels.MaternKarhunenLoeveKernel`.
737 """
738 eigensystem = self._get_eigensystem(num)
739 idx = (
740 eigensystem[f"{hodge_type}_idx"]
741 if hodge_type is not None
742 else list(range(num))
743 )
744 eigenfunctions = EigenfunctionsFromEigenvectors(
745 take_along_axis(
746 eigensystem["evecs"],
747 B.cast(int_like(eigensystem["evecs"]), np.array(idx)[None, :]),
748 axis=-1,
749 ),
750 index_from_one=True,
751 )
752 return eigenfunctions
754 def get_eigenvectors(self, num: int, hodge_type: Optional[str] = None) -> B.Numeric:
755 """
756 Eigenvectors of the Laplacian corresponding to the first `num` levels
757 (i.e., in this case, `num` eigenpairs). If `type` is specified, returns
758 only the eigenvectors of that type.
760 :param num:
761 Number of eigenvectors to return.
763 :param hodge_type:
764 The type of the eigenpairs. It can be "harmonic", "gradient", or "curl".
766 :return:
767 (n, 1)-shaped array containing the eigenvalues. n <= num.
768 """
769 eigensystem = self._get_eigensystem(num)
771 idx = (
772 eigensystem[f"{hodge_type}_idx"]
773 if hodge_type is not None
774 else list(range(num))
775 )
777 eigenvectors = take_along_axis(
778 eigensystem["evecs"],
779 B.cast(int_like(eigensystem["evecs"]), np.array(idx)[None, :]),
780 axis=-1,
781 )
783 return eigenvectors
785 def get_eigenvalues(self, num: int, hodge_type: Optional[str] = None) -> B.Numeric:
786 """
787 Eigenvalues of the Laplacian corresponding to the first `num` levels
788 (i.e., in this case, `num` eigenpairs). If `type` is specified, returns
789 only the eigenvalues corresponding to the eigenfunctions of that type.
791 .. warning::
792 If `type` is specified, the array can have fewer than `num` elements.
794 :param num:
795 Number of levels.
797 :param hodge_type:
798 The type of the eigenpairs. It can be "harmonic", "gradient", or "curl".
800 :return:
801 (n, 1)-shaped array containing the eigenvalues. n <= num.
802 """
803 eigensystem = self._get_eigensystem(num)
805 idx = (
806 eigensystem[f"{hodge_type}_idx"]
807 if hodge_type is not None
808 else list(range(num))
809 )
811 eigenvalues = take_along_axis(
812 eigensystem["evals"], B.cast(int_like(eigensystem["evals"]), np.array(idx))
813 )
815 return eigenvalues[:, None]
817 def get_repeated_eigenvalues(
818 self, num: int, hodge_type: Optional[str] = None
819 ) -> B.Numeric:
820 """
821 Same as :meth:`get_eigenvalues`.
823 :param num:
824 Same as :meth:`get_eigenvalues`.
825 """
826 return self.get_eigenvalues(num, hodge_type)
828 def random(self, key, number):
829 key, random_edges = B.randint(
830 key,
831 dtype_integer(key),
832 number,
833 1,
834 lower=1,
835 upper=self.num_edges + 1,
836 )
838 return key, random_edges
840 @property
841 def element_shape(self):
842 """
843 :return:
844 [1].
845 """
846 return [1]
848 @property
849 def element_dtype(self):
850 """
851 :return:
852 B.Int.
853 """
854 return B.Int