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

1"""Least-squares solver with NaN-aware row filtering.""" 

2 

3from __future__ import annotations 

4 

5import numpy as np 

6import numpy.typing as npt 

7 

8from ..core.exceptions import DEFAULT_COND_THRESHOLD, DimensionMismatchError 

9from ..core.exceptions import warn_ill_conditioned as _warn_ill_conditioned 

10from ..core.types import Matrix, Vector 

11 

12 

13def _condition_number(sv: npt.NDArray[np.floating]) -> float: 

14 """Condition number ``sv[0] / sv[-1]`` from descending singular values. 

15 

16 Returns ``inf`` when the smallest singular value is zero and ``1.0`` when 

17 there are no singular values (an empty valid sub-matrix). 

18 """ 

19 if sv.size == 0: 

20 return 1.0 

21 if sv[-1] > 0: 

22 return float(sv[0] / sv[-1]) 

23 return float("inf") 

24 

25 

26def lstsq( 

27 matrix: Matrix, 

28 rhs: Vector, 

29 cond_threshold: float = DEFAULT_COND_THRESHOLD, 

30) -> tuple[Vector, Vector, int, Vector]: 

31 """Solve an overdetermined or underdetermined system in the least-squares sense. 

32 

33 Rows where any entry in *matrix* or the corresponding entry in *rhs* is 

34 non-finite are excluded before solving. The returned solution vector 

35 always has length equal to the number of columns in *matrix*. When the 

36 effective condition number of the valid sub-matrix exceeds 

37 *cond_threshold*, an ``IllConditionedMatrixWarning`` is emitted. 

38 

39 Args: 

40 matrix: Coefficient matrix of shape ``(m, n)``. 

41 rhs: Right-hand side vector of length ``m``. 

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

43 emitted. Defaults to ``1e12``. 

44 

45 Returns: 

46 A four-tuple ``(x, residuals, rank, sv)`` matching the convention of 

47 :func:`numpy.linalg.lstsq`: 

48 

49 - ``x`` — least-squares solution of shape ``(n,)``. 

50 - ``residuals`` — sum of squared residuals; empty when the solution is 

51 not unique or all rows are invalid. 

52 - ``rank`` — effective rank of the valid sub-matrix. 

53 - ``sv`` — singular values of the valid sub-matrix in descending order. 

54 

55 Raises: 

56 DimensionMismatchError: If ``rhs`` length does not match the number of 

57 rows in *matrix*. 

58 

59 Example: 

60 >>> import numpy as np 

61 >>> from cvx.linalg import lstsq 

62 >>> A = np.array([[1.0, 1.0], [1.0, 2.0], [1.0, 3.0]]) 

63 >>> b = np.array([6.0, 5.0, 7.0]) 

64 >>> x, res, rank, sv = lstsq(A, b) 

65 >>> int(rank) 

66 2 

67 

68 NaN rows are silently dropped: 

69 

70 >>> A_nan = np.array([[1.0, 1.0], [np.nan, 2.0], [1.0, 3.0]]) 

71 >>> b_nan = np.array([6.0, 5.0, 7.0]) 

72 >>> x2, _, rank2, _ = lstsq(A_nan, b_nan) 

73 >>> int(rank2) 

74 2 

75 """ 

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

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

78 

79 n_cols = matrix.shape[1] 

80 

81 # Filter rows that contain any non-finite value in matrix or rhs. 

82 row_mask = np.isfinite(matrix).all(axis=1) & np.isfinite(rhs) 

83 sub_matrix = matrix[row_mask] 

84 sub_rhs = rhs[row_mask] 

85 

86 if sub_matrix.shape[0] == 0: 

87 return np.full(n_cols, np.nan), np.array([]), 0, np.array([]) 

88 

89 x, residuals, rank, sv = np.linalg.lstsq(sub_matrix, sub_rhs, rcond=None) 

90 

91 _warn_ill_conditioned(_condition_number(sv), cond_threshold) 

92 

93 return ( 

94 x.astype(np.float64, copy=False), 

95 residuals.astype(np.float64, copy=False), 

96 int(rank), 

97 sv.astype(np.float64, copy=False), 

98 )