Coverage for src / cvx / linalg / 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 .cholesky import cholesky as _cholesky 

8from .exceptions import ( 

9 DimensionMismatchError, 

10 NonSquareMatrixError, 

11 SingularMatrixError, 

12) 

13from .exceptions import ( 

14 check_and_warn_condition as _check_and_warn_condition, 

15) 

16from .valid import valid 

17 

18_DEFAULT_COND_THRESHOLD: float = 1e12 

19"""Default condition-number threshold above which a warning is emitted.""" 

20 

21 

22def solve( 

23 matrix: np.ndarray, 

24 rhs: np.ndarray, 

25 cond_threshold: float = _DEFAULT_COND_THRESHOLD, 

26) -> np.ndarray: 

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

28 

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

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

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

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

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

34 ``IllConditionedMatrixWarning`` is emitted. 

35 

36 Args: 

37 matrix: Square coefficient matrix. 

38 rhs: Right-hand side vector. 

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

40 emitted. Defaults to ``1e12``. 

41 

42 Returns: 

43 A solution vector 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 ``rhs`` size does not match the matrix dimension. 

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

50 

51 Example: 

52 >>> import numpy as np 

53 >>> from cvx.linalg import solve 

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

55 [1.0, 2.0] 

56 

57 NaN-masked entries are skipped: 

58 

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

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

61 [2.0, nan] 

62 """ 

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

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

65 

66 if rhs.size != matrix.shape[0]: 

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

68 

69 solution = np.nan * np.ones(rhs.size) 

70 mask, submatrix = valid(matrix) 

71 

72 if mask.any(): 

73 _check_and_warn_condition(submatrix, cond_threshold) 

74 try: 

75 solution[mask] = _cholesky(submatrix, rhs[mask]) 

76 except np.linalg.LinAlgError as exc: 

77 raise SingularMatrixError(str(exc)) from exc 

78 

79 return solution