Coverage for suppy\utils\_array_helper.py: 61%

1"""General class for easier matrix operations.""" 

2import numpy as np 

3from scipy import sparse 

4import warnings 

5 

6try: 

7 import cupy as cp 

8 import cupyx.scipy.sparse as csparse 

9 import cupyx.scipy.sparse.linalg as clinalg 

10 

11 NO_GPU = False 

12 

13except ImportError: 

14 NO_GPU = True 

15 cp = None 

16 csparse = None 

17 

18 

19class LinearMapping: 

20 """This class is used to allow interoperatibility between numpy, scipy etc.""" 

21 

22 def __init__(self, A): 

23 self.gpu = False 

24 

25 if not NO_GPU: 

26 if isinstance(A, cp.ndarray): 

27 self.flag = "cupy_full" 

28 self.gpu = True 

29 if A.ndim == 1: 

30 self.A = cp.array([A]) 

31 elif A.ndim > 2: 

32 raise ValueError("A must be a 2D array or a sparse matrix.") 

33 else: 

34 self.A = A 

35 return 

36 

37 elif csparse.issparse(A): 

38 self.flag = "cupy_sparse" 

39 self.gpu = True 

40 if isinstance(A, csparse.csr_matrix): 

41 self.A = A 

42 else: 

43 warnings.warn( 

44 "Converting to csr matrix format. This creates a copy of your matrix!", 

45 stacklevel=2, 

46 ) 

47 self.A = csparse.csr_matrix(A) 

48 return 

49 

50 if isinstance(A, np.ndarray): 

51 self.flag = "numpy" 

52 if A.ndim == 1: 

53 self.A = np.array([A]) 

54 elif A.ndim > 2: 

55 raise ValueError("A must be a 2D array or a sparse matrix.") 

56 else: 

57 self.A = A 

58 

59 elif sparse.issparse(A): 

60 self.flag = "scipy_sparse" 

61 if isinstance(A, sparse.csr_array): 

62 self.A = A 

63 else: 

64 warnings.warn( 

65 "Converting to csr array format. This creates a copy of your array!", 

66 stacklevel=2, 

67 ) 

68 self.A = sparse.csr_array(A) 

69 

70 else: 

71 raise ValueError( 

72 f"Unsupported matrix type: {type(A)}. Expected np.ndarray, scipy sparse, " 

73 "or (if CuPy is available) cp.ndarray or cupy sparse." 

74 ) 

75 

76 @staticmethod 

77 def get_flags(A): 

78 _gpu = False 

79 _flag = None 

80 if not NO_GPU: 

81 if isinstance(A, cp.ndarray): 

82 return "cupy_full", True 

83 elif csparse.issparse(A): 

84 return "cupy_sparse", True 

85 

86 if isinstance(A, np.ndarray): 

87 _flag = "numpy" 

88 elif sparse.issparse(A): 

89 _flag = "scipy_sparse" 

90 return _flag, _gpu 

91 

92 # Representation 

93 def __str__(self): 

94 return self.A.__str__() 

95 

96 def __repr__(self): 

97 return self.A.__repr__() 

98 

99 # Attribute access 

100 

101 def __getattr__(self, attr): 

102 return getattr(self.A, attr) 

103 

104 # Get and set elements 

105 def __getitem__(self, key): 

106 return self.A[key] 

107 

108 def __setitem__(self, key, value): 

109 self.A[key] = value 

110 

111 def __eq__(self, other): 

112 return self.A == other 

113 

114 # Mathematical operators 

115 

116 def __add__(self, other): 

117 return self.A + other 

118 

119 def __radd__(self, other): 

120 return other + self.A 

121 

122 def __sub__(self, other): 

123 return self.A - other 

124 

125 def __rsub__(self, other): 

126 return other - self.A 

127 

128 def __mul__(self, other): 

129 return self.A * other 

130 

131 def __rmul__(self, other): 

132 return other * self.A 

133 

134 def __truediv__(self, other): 

135 return self.A / other 

136 

137 def __rtruediv__(self, other): 

138 return other / self.A 

139 

140 def __pow__(self, other): 

141 return self.A**other 

142 

143 def __matmul__(self, other): 

144 return self.A @ other 

145 

146 def __rmatmul__(self, other): 

147 return other @ self.A 

148 

149 def __iter__(self): 

150 return self.A.__iter__() 

151 

152 def __len__(self): 

153 return self.A.__len__() 

154 

155 def get_norm(self, order=None, power=1): 

156 """Get the norm of the matrix.""" 

157 if self.flag == "numpy": 

158 return np.linalg.norm(self.A, ord=order) ** power 

159 

160 if self.flag == "scipy_sparse": 

161 return sparse.linalg.norm(self.A, ord=order) ** power 

162 

163 if self.flag == "cupy_full": 

164 return cp.linalg.norm(self.A, ord=order) ** power 

165 

166 if self.flag == "cupy_sparse": 

167 return clinalg.norm(self.A, ord=order) ** power 

168 raise ValueError("Unknown flag.") 

169 

170 def row_norm(self, order=None, power=1): 

171 """Get the row norms of the matrix.""" 

172 if self.flag == "numpy": 

173 return np.linalg.norm(self.A, axis=1, ord=order) ** power 

174 

175 if self.flag == "scipy_sparse": 

176 return sparse.linalg.norm(self.A, axis=1, ord=order) ** power 

177 

178 if self.flag == "cupy_full": 

179 return cp.linalg.norm(self.A, axis=1, ord=order) ** power 

180 

181 if self.flag == "cupy_sparse": 

182 return clinalg.norm(self.A, axis=1, ord=order) ** power 

183 

184 raise ValueError("Unknown flag.") 

185 

186 def single_map(self, x, i): 

187 """Apply a linear map to a single row of the matrix.""" 

188 if self.flag in ["numpy", "cupy_full"]: 

189 return self.A[i] @ x 

190 

191 if self.flag in ["scipy_sparse", "cupy_sparse"]: 

192 idx1, idx2 = self.A.indptr[i], self.A.indptr[i + 1] 

193 return self.A.data[idx1:idx2] @ x[self.A.indices[idx1:idx2]] 

194 raise ValueError("Unknown flag.") 

195 

196 def index_map(self, x, idx): 

197 """Apply a linear map to a subset of the matrix.""" 

198 return self.A[idx] @ x 

199 

200 def update_step(self, x, c, i): 

201 """Add a scaled row of the matrix to x in-place: x += A[i] * c.""" 

202 if self.flag in ["numpy", "cupy_full"]: 

203 x += self.A[i] * c 

204 

205 elif self.flag in ["scipy_sparse", "cupy_sparse"]: 

206 idx1, idx2 = self.A.indptr[i], self.A.indptr[i + 1] 

207 x[self.A.indices[idx1:idx2]] += self.A.data[idx1:idx2] * c 

208 else: 

209 raise ValueError("Unknown flag.") 

210 

211 def getrow(self, i): 

212 """Get a row of the matrix.""" 

213 if self.flag in ["numpy", "cupy_full"]: 

214 return self.A[i] 

215 

216 if self.flag in ["scipy_sparse", "cupy_sparse"]: 

217 return self.A.getrow(i) 

218 

219 raise ValueError("Unknown flag.")