Coverage for src/cvx/linalg/decomposition/power_iteration.py: 100%

1"""Dominant eigenpair estimation via power iteration.""" 

2 

3from __future__ import annotations 

4 

5from collections.abc import Callable 

6from typing import TYPE_CHECKING 

7 

8import numpy as np 

9 

10from ..core.exceptions import NonSquareMatrixError, NotAMatrixError 

11from ..core.types import Matrix, Vector 

12 

13if TYPE_CHECKING: 

14 from ..operators import SymmetricOperator 

15 

16_CALLABLE_NEEDS_N_MESSAGE = "power_iteration needs `n` when `operator` is a bare callable" 

17 

18 

19def power_iteration( 

20 operator: Matrix | SymmetricOperator | Callable[[Vector], Vector], 

21 *, 

22 n: int | None = None, 

23 n_iter: int = 1000, 

24 tol: float = 1e-9, 

25 seed: int | None = None, 

26) -> tuple[float, Vector]: 

27 """Estimate the dominant eigenpair of a real symmetric operator. 

28 

29 Repeatedly applies *operator* to a random unit vector, renormalizing each 

30 step, until the Rayleigh-quotient eigenvalue estimate stops changing. The 

31 iterate converges to the eigenvector whose eigenvalue is largest in 

32 magnitude; the matching eigenvalue (returned with its sign) comes from the 

33 Rayleigh quotient ``v.T @ (operator @ v)``. 

34 

35 *operator* may be given three ways, so the leading eigenvalue can be 

36 estimated **matrix-free** (e.g. the Lipschitz constant of a gradient step): 

37 

38 * a dense symmetric ``(n, n)`` array; 

39 * any :class:`~cvx.linalg.SymmetricOperator` (its :meth:`matvec` and ``n`` 

40 drive the iteration, so no ``n x n`` matrix is formed); or 

41 * a callable ``v -> A @ v`` together with the dimension *n*. 

42 

43 Each iteration costs a single application, so this is far cheaper than a full 

44 :func:`~cvx.linalg.eigh` when only the leading eigenpair is needed and there 

45 is a clear spectral gap. Symmetry is assumed; only the result's 

46 interpretation as an eigenpair relies on it. Like :func:`~cvx.linalg.svd`, 

47 this is a raw primitive and is **not** NaN-aware. 

48 

49 Args: 

50 operator: A symmetric matrix, a :class:`~cvx.linalg.SymmetricOperator`, 

51 or a callable applying it to a vector. 

52 n: Dimension of the operator. Required when *operator* is a bare callable; 

53 ignored otherwise (taken from the array shape or the operator's ``n``). 

54 n_iter: Maximum number of iterations. Defaults to 1000. 

55 tol: Convergence tolerance on the relative change of the eigenvalue 

56 estimate between iterations. Defaults to ``1e-9``. 

57 seed: Seed for the random starting vector, for reproducibility. If 

58 ``None``, uses the current NumPy random state. 

59 

60 Returns: 

61 Tuple ``(eigenvalue, eigenvector)`` where *eigenvalue* is the signed 

62 dominant eigenvalue estimate (a ``float``) and *eigenvector* is the 

63 corresponding unit-norm eigenvector. The eigenvector sign is arbitrary. 

64 

65 Raises: 

66 NotAMatrixError: If *operator* is an array that is not 2-D. 

67 NonSquareMatrixError: If *operator* is an array that is not square. 

68 ValueError: If *operator* is a bare callable and *n* is not given. 

69 

70 Example: 

71 >>> import numpy as np 

72 >>> from cvx.linalg import power_iteration 

73 >>> matrix = np.diag([3.0, 2.0, 1.0]) 

74 >>> eigenvalue, eigenvector = power_iteration(matrix, seed=0) 

75 >>> bool(np.isclose(eigenvalue, 3.0)) 

76 True 

77 >>> bool(np.isclose(abs(eigenvector[0]), 1.0)) 

78 True 

79 

80 It runs matrix-free on a :class:`~cvx.linalg.SymmetricOperator`, so the 

81 leading eigenvalue of ``M.T @ M`` needs no ``n x n`` matrix: 

82 

83 >>> from cvx.linalg import GramOperator 

84 >>> rng = np.random.default_rng(0) 

85 >>> M = rng.standard_normal((20, 5)) 

86 >>> lam, _ = power_iteration(GramOperator(M), seed=0) 

87 >>> bool(np.isclose(lam, np.linalg.eigvalsh(M.T @ M)[-1])) 

88 True 

89 """ 

90 apply, dim = _resolve(operator, n) 

91 

92 rng = np.random.default_rng(seed) 

93 v = rng.standard_normal((dim,)) 

94 v = v / np.linalg.norm(v) 

95 

96 eigenvalue = float(v @ apply(v)) 

97 for _ in range(n_iter): 

98 w = apply(v) 

99 norm = float(np.linalg.norm(w)) 

100 if norm < 1e-15: 

101 # operator annihilates the iterate: the dominant eigenvalue is zero. 

102 return 0.0, v 

103 v = w / norm 

104 new_eigenvalue = float(v @ apply(v)) 

105 if abs(new_eigenvalue - eigenvalue) <= tol * max(1.0, abs(new_eigenvalue)): 

106 return new_eigenvalue, v 

107 eigenvalue = new_eigenvalue 

108 

109 return eigenvalue, v 

110 

111 

112def _resolve( 

113 operator: Matrix | SymmetricOperator | Callable[[Vector], Vector], 

114 n: int | None, 

115) -> tuple[Callable[[Vector], Vector], int]: 

116 """Return ``(apply, dimension)`` for an array, a SymmetricOperator, or a callable.""" 

117 # Imported lazily: operators depend on decomposition.cholesky, so a module-level 

118 # import would close an import cycle (hence the TYPE_CHECKING-only import above). 

119 from ..operators import SymmetricOperator 

120 

121 if isinstance(operator, SymmetricOperator): 

122 return operator.matvec, operator.n 

123 if not isinstance(operator, np.ndarray) and callable(operator): 

124 if n is None: 

125 raise ValueError(_CALLABLE_NEEDS_N_MESSAGE) 

126 return operator, int(n) 

127 array = np.asarray(operator) 

128 if array.ndim != 2: 

129 raise NotAMatrixError(array.ndim, func="power_iteration") 

130 rows, cols = array.shape 

131 if rows != cols: 

132 raise NonSquareMatrixError(rows, cols) 

133 return lambda v: array @ v, cols