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

1"""Weighted sum of symmetric operators.""" 

2 

3from __future__ import annotations 

4 

5from collections.abc import Sequence 

6 

7from ..core.exceptions import DimensionMismatchError 

8from ..core.types import Matrix, Vector 

9from .base import SymmetricOperator 

10 

11_EMPTY_TERMS_MESSAGE = "SumOperator needs at least one term" 

12_NO_SOLVE_MESSAGE = "SumOperator has no direct free-block solve; run a Krylov method (e.g. CG) over apply_free" 

13_NO_RCOND_MESSAGE = "SumOperator has no structural rcond_free estimate" 

14 

15 

16class SumOperator(SymmetricOperator): 

17 """A weighted sum of symmetric operators, ``A = sum_i coeff_i * op_i``. 

18 

19 Composes several :class:`SymmetricOperator` terms that act on the same 

20 dimension into one operator whose products are the (weighted) sums of the 

21 terms' products. This expresses forms a single backend cannot -- e.g. a 

22 data term plus a structured target, 

23 ``(1 - alpha) * GramOperator(M) + alpha * FactorOperator(...)`` -- while 

24 every product stays matrix-free (each term applies its own). 

25 

26 **Forward-only.** A general sum has no cheap direct free-block solve or 

27 structural conditioning estimate, so :meth:`solve_free` and 

28 :meth:`rcond_free` raise :class:`NotImplementedError`. Feed 

29 :meth:`apply_free` to a Krylov solver (conjugate gradients) instead -- which 

30 is exactly how such a composite is used. :attr:`diag` *is* available (the 

31 weighted sum of the terms' diagonals, provided every term exposes one) and 

32 supplies the Jacobi preconditioner for that solve. 

33 

34 Args: 

35 terms: A non-empty sequence of ``(coeff, operator)`` pairs. All operators 

36 must share the same dimension ``n``. 

37 

38 Example: 

39 >>> import numpy as np 

40 >>> from cvx.linalg import DenseOperator, SumOperator 

41 >>> P = DenseOperator(np.array([[2.0, 0.0], [0.0, 1.0]])) 

42 >>> Q = DenseOperator(np.array([[1.0, 1.0], [1.0, 1.0]])) 

43 >>> op = SumOperator([(0.5, P), (2.0, Q)]) 

44 >>> A = 0.5 * np.array([[2.0, 0.0], [0.0, 1.0]]) + 2.0 * np.ones((2, 2)) 

45 >>> np.allclose(op.matvec(np.array([1.0, -1.0])), A @ np.array([1.0, -1.0])) 

46 True 

47 >>> free = np.array([0, 1]) 

48 >>> np.allclose(op.apply_free(free, np.array([1.0, 1.0])), A @ np.ones(2)) 

49 True 

50 >>> np.allclose(op.diag, np.diag(A)) # Jacobi preconditioner 

51 True 

52 """ 

53 

54 def __init__(self, terms: Sequence[tuple[float, SymmetricOperator]]) -> None: 

55 """Store the ``(coeff, operator)`` terms after checking they share a dimension.""" 

56 terms = list(terms) 

57 if not terms: 

58 raise ValueError(_EMPTY_TERMS_MESSAGE) 

59 dims = {op.n for _, op in terms} 

60 if len(dims) != 1: 

61 raise DimensionMismatchError(min(dims), max(dims)) 

62 self._terms: list[tuple[float, SymmetricOperator]] = [(float(c), op) for c, op in terms] 

63 self._n = dims.pop() 

64 

65 @property 

66 def n(self) -> int: 

67 """Dimension shared by the terms.""" 

68 return int(self._n) 

69 

70 @property 

71 def diag(self) -> Vector: 

72 """The weighted sum of the terms' diagonals. 

73 

74 Raises: 

75 NotImplementedError: If any term does not expose its diagonal. 

76 """ 

77 coeff, op = self._terms[0] 

78 result: Vector = coeff * op.diag 

79 for coeff, op in self._terms[1:]: 

80 result = result + coeff * op.diag 

81 return result 

82 

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

84 """Return ``A @ x = sum_i coeff_i * op_i.matvec(x)``.""" 

85 coeff, op = self._terms[0] 

86 result = coeff * op.matvec(x) 

87 for coeff, op in self._terms[1:]: 

88 result = result + coeff * op.matvec(x) 

89 return result 

90 

91 def restricted(self, free: object) -> SumOperator: 

92 """Return the sum of the terms' restricted operators (same coefficients).""" 

93 return SumOperator([(c, op.restricted(free)) for c, op in self._terms]) 

94 

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

96 """Return ``A[rows, cols] @ v`` as the weighted sum of the terms' sub-block products.""" 

97 coeff, op = self._terms[0] 

98 result = coeff * op.block_matvec(rows, cols, v) 

99 for coeff, op in self._terms[1:]: 

100 result = result + coeff * op.block_matvec(rows, cols, v) 

101 return result 

102 

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

104 """Not available: a sum has no cheap direct free-block solve. 

105 

106 Raises: 

107 NotImplementedError: Always. Run a Krylov solver over :meth:`apply_free`. 

108 """ 

109 raise NotImplementedError(_NO_SOLVE_MESSAGE) 

110 

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

112 """Not available: a sum has no structural conditioning estimate. 

113 

114 Raises: 

115 NotImplementedError: Always. Estimate conditioning from the assembled 

116 free block or a Krylov method if needed. 

117 """ 

118 raise NotImplementedError(_NO_RCOND_MESSAGE)