Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.9/dist-packages/scipy/_lib/array_api_compat/common/_linalg.py: 40%

1from __future__ import annotations 

2 

3from typing import TYPE_CHECKING, NamedTuple 

4if TYPE_CHECKING: 

5 from typing import Literal, Optional, Sequence, Tuple, Union 

6 from ._typing import ndarray 

7 

8import numpy as np 

9if np.__version__[0] == "2": 

10 from numpy.lib.array_utils import normalize_axis_tuple 

11else: 

12 from numpy.core.numeric import normalize_axis_tuple 

13 

14from ._aliases import matmul, matrix_transpose, tensordot, vecdot, isdtype 

15from .._internal import get_xp 

16 

17# These are in the main NumPy namespace but not in numpy.linalg 

18def cross(x1: ndarray, x2: ndarray, /, xp, *, axis: int = -1, **kwargs) -> ndarray: 

19 return xp.cross(x1, x2, axis=axis, **kwargs) 

20 

21def outer(x1: ndarray, x2: ndarray, /, xp, **kwargs) -> ndarray: 

22 return xp.outer(x1, x2, **kwargs) 

23 

24class EighResult(NamedTuple): 

25 eigenvalues: ndarray 

26 eigenvectors: ndarray 

27 

28class QRResult(NamedTuple): 

29 Q: ndarray 

30 R: ndarray 

31 

32class SlogdetResult(NamedTuple): 

33 sign: ndarray 

34 logabsdet: ndarray 

35 

36class SVDResult(NamedTuple): 

37 U: ndarray 

38 S: ndarray 

39 Vh: ndarray 

40 

41# These functions are the same as their NumPy counterparts except they return 

42# a namedtuple. 

43def eigh(x: ndarray, /, xp, **kwargs) -> EighResult: 

44 return EighResult(*xp.linalg.eigh(x, **kwargs)) 

45 

46def qr(x: ndarray, /, xp, *, mode: Literal['reduced', 'complete'] = 'reduced', 

47 **kwargs) -> QRResult: 

48 return QRResult(*xp.linalg.qr(x, mode=mode, **kwargs)) 

49 

50def slogdet(x: ndarray, /, xp, **kwargs) -> SlogdetResult: 

51 return SlogdetResult(*xp.linalg.slogdet(x, **kwargs)) 

52 

53def svd(x: ndarray, /, xp, *, full_matrices: bool = True, **kwargs) -> SVDResult: 

54 return SVDResult(*xp.linalg.svd(x, full_matrices=full_matrices, **kwargs)) 

55 

56# These functions have additional keyword arguments 

57 

58# The upper keyword argument is new from NumPy 

59def cholesky(x: ndarray, /, xp, *, upper: bool = False, **kwargs) -> ndarray: 

60 L = xp.linalg.cholesky(x, **kwargs) 

61 if upper: 

62 U = get_xp(xp)(matrix_transpose)(L) 

63 if get_xp(xp)(isdtype)(U.dtype, 'complex floating'): 

64 U = xp.conj(U) 

65 return U 

66 return L 

67 

68# The rtol keyword argument of matrix_rank() and pinv() is new from NumPy. 

69# Note that it has a different semantic meaning from tol and rcond. 

70def matrix_rank(x: ndarray, 

71 /, 

72 xp, 

73 *, 

74 rtol: Optional[Union[float, ndarray]] = None, 

75 **kwargs) -> ndarray: 

76 # this is different from xp.linalg.matrix_rank, which supports 1 

77 # dimensional arrays. 

78 if x.ndim < 2: 

79 raise xp.linalg.LinAlgError("1-dimensional array given. Array must be at least two-dimensional") 

80 S = xp.linalg.svd(x, compute_uv=False, **kwargs) 

81 if rtol is None: 

82 tol = S.max(axis=-1, keepdims=True) * max(x.shape[-2:]) * xp.finfo(S.dtype).eps 

83 else: 

84 # this is different from xp.linalg.matrix_rank, which does not 

85 # multiply the tolerance by the largest singular value. 

86 tol = S.max(axis=-1, keepdims=True)*xp.asarray(rtol)[..., xp.newaxis] 

87 return xp.count_nonzero(S > tol, axis=-1) 

88 

89def pinv(x: ndarray, /, xp, *, rtol: Optional[Union[float, ndarray]] = None, **kwargs) -> ndarray: 

90 # this is different from xp.linalg.pinv, which does not multiply the 

91 # default tolerance by max(M, N). 

92 if rtol is None: 

93 rtol = max(x.shape[-2:]) * xp.finfo(x.dtype).eps 

94 return xp.linalg.pinv(x, rcond=rtol, **kwargs) 

95 

96# These functions are new in the array API spec 

97 

98def matrix_norm(x: ndarray, /, xp, *, keepdims: bool = False, ord: Optional[Union[int, float, Literal['fro', 'nuc']]] = 'fro') -> ndarray: 

99 return xp.linalg.norm(x, axis=(-2, -1), keepdims=keepdims, ord=ord) 

100 

101# svdvals is not in NumPy (but it is in SciPy). It is equivalent to 

102# xp.linalg.svd(compute_uv=False). 

103def svdvals(x: ndarray, /, xp) -> Union[ndarray, Tuple[ndarray, ...]]: 

104 return xp.linalg.svd(x, compute_uv=False) 

105 

106def vector_norm(x: ndarray, /, xp, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keepdims: bool = False, ord: Optional[Union[int, float]] = 2) -> ndarray: 

107 # xp.linalg.norm tries to do a matrix norm whenever axis is a 2-tuple or 

108 # when axis=None and the input is 2-D, so to force a vector norm, we make 

109 # it so the input is 1-D (for axis=None), or reshape so that norm is done 

110 # on a single dimension. 

111 if axis is None: 

112 # Note: xp.linalg.norm() doesn't handle 0-D arrays 

113 x = x.ravel() 

114 _axis = 0 

115 elif isinstance(axis, tuple): 

116 # Note: The axis argument supports any number of axes, whereas 

117 # xp.linalg.norm() only supports a single axis for vector norm. 

118 normalized_axis = normalize_axis_tuple(axis, x.ndim) 

119 rest = tuple(i for i in range(x.ndim) if i not in normalized_axis) 

120 newshape = axis + rest 

121 x = xp.transpose(x, newshape).reshape( 

122 (xp.prod([x.shape[i] for i in axis], dtype=int), *[x.shape[i] for i in rest])) 

123 _axis = 0 

124 else: 

125 _axis = axis 

126 

127 res = xp.linalg.norm(x, axis=_axis, ord=ord) 

128 

129 if keepdims: 

130 # We can't reuse xp.linalg.norm(keepdims) because of the reshape hacks 

131 # above to avoid matrix norm logic. 

132 shape = list(x.shape) 

133 _axis = normalize_axis_tuple(range(x.ndim) if axis is None else axis, x.ndim) 

134 for i in _axis: 

135 shape[i] = 1 

136 res = xp.reshape(res, tuple(shape)) 

137 

138 return res 

139 

140# xp.diagonal and xp.trace operate on the first two axes whereas these 

141# operates on the last two 

142 

143def diagonal(x: ndarray, /, xp, *, offset: int = 0, **kwargs) -> ndarray: 

144 return xp.diagonal(x, offset=offset, axis1=-2, axis2=-1, **kwargs) 

145 

146def trace(x: ndarray, /, xp, *, offset: int = 0, dtype=None, **kwargs) -> ndarray: 

147 if dtype is None: 

148 if x.dtype == xp.float32: 

149 dtype = xp.float64 

150 elif x.dtype == xp.complex64: 

151 dtype = xp.complex128 

152 return xp.asarray(xp.trace(x, offset=offset, dtype=dtype, axis1=-2, axis2=-1, **kwargs)) 

153 

154__all__ = ['cross', 'matmul', 'outer', 'tensordot', 'EighResult', 

155 'QRResult', 'SlogdetResult', 'SVDResult', 'eigh', 'qr', 'slogdet', 

156 'svd', 'cholesky', 'matrix_rank', 'pinv', 'matrix_norm', 

157 'matrix_transpose', 'svdvals', 'vecdot', 'vector_norm', 'diagonal', 

158 'trace']