Coverage for src/cvx/linalg/solve/solve.py: 100%

1"""Linear-system helpers that ignore matrix rows and columns with non-finite diagonals.""" 

2 

3from __future__ import annotations 

4 

5import numpy as np 

6 

7from ..core.condition import DEFAULT_COND_THRESHOLD 

8from ..core.condition import ( 

9 check_and_warn_condition as _check_and_warn_condition, 

10) 

11from ..core.exceptions import ( 

12 DimensionMismatchError, 

13 NonSquareMatrixError, 

14 SingularMatrixError, 

15) 

16from ..core.types import Matrix, Vector 

17from ..core.valid import valid 

18from ..decomposition.cholesky import cholesky_solve as _cholesky_solve 

19 

20 

21def solve( 

22 matrix: Matrix, 

23 rhs: Vector | Matrix, 

24 cond_threshold: float | None = DEFAULT_COND_THRESHOLD, 

25) -> Vector | Matrix: 

26 """Solve a linear system restricted to the valid submatrix. 

27 

28 Rows and columns with non-finite diagonal entries are excluded from the 

29 solve; the corresponding positions in the result are set to NaN. Cholesky 

30 decomposition is attempted first for numerical stability and falls back to 

31 LU decomposition for non-positive-definite matrices. When the condition 

32 number of the valid sub-matrix exceeds *cond_threshold*, an 

33 ``IllConditionedMatrixWarning`` is emitted. 

34 

35 Args: 

36 matrix: Square coefficient matrix of shape ``(n, n)``. 

37 rhs: Right-hand side vector of length ``n`` or matrix of shape ``(n, k)``. 

38 cond_threshold: Condition-number threshold above which a warning is 

39 emitted. Defaults to ``1e12``. ``None`` skips the check, and the 

40 SVD that computes the condition number, entirely. 

41 

42 Returns: 

43 A solution array with the same shape as ``rhs``. Entries mapped to 

44 invalid rows or columns are returned as ``NaN``. 

45 

46 Raises: 

47 NonSquareMatrixError: If the matrix is not square. 

48 DimensionMismatchError: If the leading dimension of ``rhs`` does not 

49 match the matrix dimension. 

50 SingularMatrixError: If the valid sub-matrix is singular. 

51 

52 Example: 

53 >>> import numpy as np 

54 >>> from cvx.linalg import solve 

55 >>> solve(np.eye(2), np.array([1.0, 2.0])).tolist() 

56 [1.0, 2.0] 

57 

58 NaN-masked entries are skipped: 

59 

60 >>> matrix = np.array([[4.0, 0.0], [0.0, np.nan]]) 

61 >>> solve(matrix, np.array([8.0, 1.0])).tolist() 

62 [2.0, nan] 

63 

64 Matrix right-hand sides are supported: 

65 

66 >>> solve(np.eye(2), np.array([[1.0, 2.0], [3.0, 4.0]])).tolist() 

67 [[1.0, 2.0], [3.0, 4.0]] 

68 """ 

69 if matrix.shape[0] != matrix.shape[1]: 

70 raise NonSquareMatrixError(matrix.shape[0], matrix.shape[1]) 

71 

72 if rhs.shape[0] != matrix.shape[0]: 

73 raise DimensionMismatchError(rhs.shape[0], matrix.shape[0]) 

74 

75 solution = np.full(rhs.shape, np.nan) 

76 mask, submatrix = valid(matrix) 

77 

78 if mask.any(): 

79 _check_and_warn_condition(submatrix, cond_threshold) 

80 try: 

81 solution[mask] = _cholesky_solve(submatrix, rhs[mask]) 

82 except np.linalg.LinAlgError as exc: 

83 raise SingularMatrixError(str(exc)) from exc 

84 

85 return solution