Coverage for src/cvx/linalg/operators/factor.py: 100%

1""":class:`FactorOperator`: diagonal-plus-low-rank ``A = diag(d) + U @ Delta @ U.T``.""" 

2 

3from __future__ import annotations 

4 

5import numpy as np 

6 

7from ..core.exceptions import DimensionMismatchError, NonSquareMatrixError, NotAMatrixError 

8from ..core.types import Matrix, Vector 

9from ..decomposition.cholesky import cholesky_solve 

10from .base import SymmetricOperator, as_index 

11 

12_DIAGONAL_NDIM_MESSAGE = "diagonal must be a 1-D array" 

13_DIAGONAL_POSITIVE_MESSAGE = "diagonal entries must be strictly positive" 

14 

15 

16def _validate_diagonal(d: Vector) -> None: 

17 """Check the diagonal is a 1-D, strictly positive vector.""" 

18 if d.ndim != 1: 

19 raise ValueError(_DIAGONAL_NDIM_MESSAGE) 

20 if np.any(d <= 0.0): 

21 raise ValueError(_DIAGONAL_POSITIVE_MESSAGE) 

22 

23 

24def _validate_loadings(u: Matrix, d: Vector) -> None: 

25 """Check the loadings are an ``n x r`` matrix whose rows match the diagonal.""" 

26 if u.ndim != 2: 

27 raise NotAMatrixError(u.ndim, func="FactorOperator") 

28 if u.shape[0] != d.shape[0]: 

29 raise DimensionMismatchError(u.shape[0], d.shape[0]) 

30 

31 

32def _validate_inner(delta: Matrix, u: Matrix) -> None: 

33 """Check the inner block is a square ``r x r`` matrix matching the loadings' rank.""" 

34 if delta.ndim != 2: 

35 raise NotAMatrixError(delta.ndim, func="FactorOperator") 

36 if delta.shape[0] != delta.shape[1]: 

37 raise NonSquareMatrixError(delta.shape[0], delta.shape[1]) 

38 if delta.shape[0] != u.shape[1]: 

39 raise DimensionMismatchError(delta.shape[0], u.shape[1]) 

40 

41 

42class FactorOperator(SymmetricOperator): 

43 """Diagonal-plus-low-rank operator ``A = diag(d) + U @ Delta @ U.T``. 

44 

45 Free-block solves use the Woodbury identity, costing ``O(len(free) r**2 + 

46 r**3)`` for a rank-``r`` factor rather than ``O(len(free)**3)``, and no 

47 ``n x n`` matrix is formed (memory ``O(n r)``). With a strictly positive 

48 diagonal *d* and positive-definite *Delta* every principal block is positive 

49 definite, so :meth:`solve_free` is always well posed. 

50 

51 Args: 

52 diagonal: The strictly positive diagonal ``d`` of length ``n``. 

53 loadings: The ``n x r`` factor loadings ``U``. 

54 inner: The ``r x r`` positive-definite inner matrix ``Delta``. 

55 

56 Example: 

57 >>> import numpy as np 

58 >>> from cvx.linalg import FactorOperator 

59 >>> d = np.array([2.0, 3.0, 4.0]) 

60 >>> U = np.array([[1.0], [0.5], [-1.0]]) 

61 >>> Delta = np.array([[2.0]]) 

62 >>> op = FactorOperator(d, U, Delta) 

63 >>> (op.n, op.k) # 3 assets, 1 factor 

64 (3, 1) 

65 >>> A = np.diag(d) + U @ Delta @ U.T 

66 >>> free, rhs = np.array([0, 2]), np.array([1.0, 1.0]) 

67 >>> np.allclose(A[np.ix_(free, free)] @ op.solve_free(free, rhs), rhs) 

68 True 

69 """ 

70 

71 def __init__(self, diagonal: Vector, loadings: Matrix, inner: Matrix) -> None: 

72 """Store the diagonal, loadings, and inner block after shape checks.""" 

73 d = np.asarray(diagonal, dtype=np.float64) 

74 u = np.asarray(loadings, dtype=np.float64) 

75 delta = np.asarray(inner, dtype=np.float64) 

76 _validate_diagonal(d) 

77 _validate_loadings(u, d) 

78 _validate_inner(delta, u) 

79 self._d = d 

80 self._u = u 

81 self._delta = delta 

82 

83 @property 

84 def n(self) -> int: 

85 """Dimension of the operator (length of the diagonal ``d``).""" 

86 return int(self._d.shape[0]) 

87 

88 @property 

89 def k(self) -> int: 

90 """Number of factors (rank ``r`` of the low-rank term; columns of ``U``).""" 

91 return int(self._u.shape[1]) 

92 

93 @property 

94 def diag(self) -> Vector: 

95 """The diagonal ``d_i + U[i] @ Delta @ U[i]``, at ``O(n r**2)`` without forming ``A``.""" 

96 result: Vector = self._d + np.einsum("ij,ij->i", self._u @ self._delta, self._u) 

97 return result 

98 

99 def matvec(self, x: Vector | Matrix) -> Vector | Matrix: 

100 """Return ``A @ x = d * x + U @ (Delta @ (U.T @ x))``.""" 

101 return (self._d * x.T).T + self._u @ (self._delta @ (self._u.T @ x)) 

102 

103 def restricted(self, free: object) -> FactorOperator: 

104 """Return ``FactorOperator(d[free], U[free], Delta)``: the free block, pre-sliced.""" 

105 free = as_index(free) 

106 return FactorOperator(self._d[free], np.ascontiguousarray(self._u[free, :]), self._delta) 

107 

108 def block_matvec(self, rows: object, cols: object, v: Vector | Matrix) -> Vector | Matrix: 

109 """Return ``A[rows, cols] @ v`` from the low-rank term and the diagonal overlap.""" 

110 rows = as_index(rows) 

111 cols = as_index(cols) 

112 low_rank = self._u[rows] @ (self._delta @ (self._u[cols].T @ v)) 

113 # Diagonal couples only positions where a row index equals a column index. 

114 common, r_idx, c_idx = np.intersect1d(rows, cols, return_indices=True) 

115 diag = np.zeros_like(low_rank) 

116 diag[r_idx] = (self._d[common] * np.asarray(v)[c_idx].T).T 

117 result: Vector | Matrix = low_rank + diag 

118 return result 

119 

120 def solve_free(self, free: object, rhs: Vector | Matrix) -> Vector | Matrix: 

121 """Solve the free block by the Woodbury identity on the ``r x r`` capacitance matrix.""" 

122 free = as_index(free) 

123 df = self._d[free] 

124 uf = self._u[free] 

125 # Woodbury: A_FF^{-1} = D^{-1} - D^{-1} U W^{-1} U.T D^{-1}, 

126 # with W = Delta^{-1} + U.T D^{-1} U. 

127 dinv_rhs = (np.asarray(rhs, dtype=np.float64).T / df).T 

128 delta_inv = np.linalg.solve(self._delta, np.eye(self._delta.shape[0])) 

129 w = delta_inv + uf.T @ ((uf.T / df).T) 

130 inner = cholesky_solve(w, uf.T @ dinv_rhs) 

131 correction = (uf @ inner).T / df 

132 result: Vector | Matrix = dinv_rhs - correction.T 

133 return result 

134 

135 def rcond_free(self, free: object) -> float: 

136 """Lower bound on the free block's reciprocal condition number, via Weyl's inequalities. 

137 

138 The free block ``diag(d_F) + U_F Delta U_F.T`` is positive definite (the 

139 positive diagonal keeps it full rank). Rather than form it, bound 

140 ``lambda_min >= min(d_F)`` and 

141 ``lambda_max <= max(d_F) + ||U_F||_2^2 * lambda_max(Delta)``; their ratio is 

142 a guaranteed lower bound on the true reciprocal condition number, at 

143 ``O(len(free) r**2 + r**3)`` and without an ``n x n`` matrix. 

144 """ 

145 free = as_index(free) 

146 if free.size == 0: 

147 return 1.0 

148 d_free = self._d[free] 

149 u_free = self._u[free] 

150 u_spectral_norm = float(np.linalg.svd(u_free, compute_uv=False)[0]) 

151 delta_max = float(np.linalg.eigvalsh(self._delta)[-1]) 

152 lam_max_upper = float(np.max(d_free)) + u_spectral_norm**2 * max(delta_max, 0.0) 

153 return float(np.min(d_free)) / lam_max_upper