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 

5from typing import cast 

6 

7import numpy as np 

8 

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

10from ..core.types import Matrix, Vector 

11from ..decomposition.cholesky import cholesky_solve 

12from .base import SymmetricOperator, as_index 

13 

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

15_DIAGONAL_POSITIVE_MESSAGE = "diagonal entries must be strictly positive" 

16 

17 

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

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

20 if d.ndim != 1: 

21 raise ValueError(_DIAGONAL_NDIM_MESSAGE) 

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

23 raise ValueError(_DIAGONAL_POSITIVE_MESSAGE) 

24 

25 

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

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

28 if u.ndim != 2: 

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

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

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

32 

33 

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

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

36 if delta.ndim != 2: 

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

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

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

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

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

42 

43 

44class FactorOperator(SymmetricOperator): 

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

46 

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

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

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

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

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

52 

53 Args: 

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

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

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

57 

58 Example: 

59 >>> import numpy as np 

60 >>> from cvx.linalg import FactorOperator 

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

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

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

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

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

66 (3, 1) 

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

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

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

70 True 

71 """ 

72 

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

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

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

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

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

78 _validate_diagonal(d) 

79 _validate_loadings(u, d) 

80 _validate_inner(delta, u) 

81 self._d = d 

82 self._u = u 

83 self._delta = delta 

84 self._delta_inv: Matrix | None = None 

85 

86 @property 

87 def n(self) -> int: 

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

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

90 

91 @property 

92 def k(self) -> int: 

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

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

95 

96 @property 

97 def diag(self) -> Vector: 

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

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

100 return result 

101 

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

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

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

105 

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

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

108 free = as_index(free) 

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

110 

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

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

113 rows = as_index(rows) 

114 cols = as_index(cols) 

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

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

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

118 diag = np.zeros_like(low_rank) 

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

120 result: Vector | Matrix = low_rank + diag 

121 return result 

122 

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

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

125 free = as_index(free) 

126 df = self._d[free] 

127 uf = self._u[free] 

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

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

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

131 w = self._inner_inverse() + uf.T @ ((uf.T / df).T) 

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

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

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

135 return result 

136 

137 def _inner_inverse(self) -> Matrix: 

138 """Return ``Delta^{-1}``, computed on the first solve and cached (``Delta`` is fixed).""" 

139 if self._delta_inv is None: 

140 self._delta_inv = cast("Matrix", np.linalg.solve(self._delta, np.eye(self._delta.shape[0]))) 

141 return self._delta_inv 

142 

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

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

145 

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

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

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

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

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

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

152 """ 

153 free = as_index(free) 

154 if free.size == 0: 

155 return 1.0 

156 d_free = self._d[free] 

157 u_free = self._u[free] 

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

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

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

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