Coverage for C:\Users\t590r\Documents\GitHub\suppy\suppy\utils\_array_helper.py: 62%

138 statements  

« prev     ^ index     » next       coverage.py v7.6.4, created at 2025-02-05 10:12 +0100

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

2import numpy as np 

3from scipy import sparse 

4 

5try: 

6 import cupy as cp 

7 import cupyx.scipy.sparse as csparse 

8 

9 NO_GPU = False 

10 

11except ImportError: 

12 NO_GPU = True 

13 cp = None 

14 csparse = None 

15 

16 

17class LinearMapping: 

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

19 

20 def __init__(self, A): 

21 self.gpu = False # flag for gpu 

22 

23 if NO_GPU is False: # is only checked when cupy is available 

24 

25 if isinstance(A, cp.ndarray): 

26 self.flag = "cupy_full" 

27 self.gpu = True # set a flag for gpu 

28 if A.ndim == 1: 

29 self.A = cp.array([A]) # wrap 1d arrays 

30 elif A.ndim > 2: 

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

32 else: 

33 self.A = A 

34 

35 elif csparse.issparse(A): 

36 self.flag = "cupy_sparse" 

37 self.gpu = True # set a flag for gpu 

38 self.A = csparse.csr_matrix(A) # transform to csr format 

39 

40 # set a flag based on class 

41 

42 if isinstance(A, np.ndarray): 

43 self.flag = "numpy" 

44 if A.ndim == 1: 

45 self.A = np.array([A]) # wrap 1d arrays 

46 

47 elif A.ndim > 2: 

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

49 

50 else: 

51 self.A = A 

52 

53 elif sparse.issparse(A): 

54 self.flag = "scipy_sparse" 

55 self.A = sparse.csr_array(A) # transform to csr format 

56 

57 @staticmethod 

58 def get_flags(A): 

59 _gpu = False 

60 _flag = None 

61 if NO_GPU is False: # is only checked when cupy is available 

62 

63 if isinstance(A, cp.ndarray): 

64 _flag = "cupy_full" 

65 _gpu = True # set a flag for gpu 

66 

67 elif csparse.issparse(A): 

68 _flag = "cupy_sparse" 

69 _gpu = True # set a flag for gpu 

70 

71 if isinstance(A, np.ndarray): 

72 _flag = "numpy" 

73 

74 elif sparse.issparse(A): 

75 _flag = "scipy_sparse" 

76 return _flag, _gpu 

77 

78 # Representation 

79 def __str__(self): 

80 return self.A.__str__() 

81 

82 def __repr__(self): 

83 return self.A.__repr__() 

84 

85 # Attribute access 

86 

87 def __getattr__(self, attr): 

88 return getattr(self.A, attr) 

89 

90 def __hasattr__(self, attr): 

91 return hasattr(self.A, attr) 

92 

93 # Get and set elements 

94 def __getitem__(self, key): 

95 return self.A[key] 

96 

97 def __setitem__(self, key, value): 

98 self.A[key] = value 

99 

100 def __eq__(self, other): 

101 return self.A == other 

102 

103 # Mathematical operators 

104 

105 def __add__(self, other): 

106 return self.A + other 

107 

108 def __radd__(self, other): 

109 return other + self.A 

110 

111 def __sub__(self, other): 

112 return self.A - other 

113 

114 def __rsub__(self, other): 

115 return other - self.A 

116 

117 def __mul__(self, other): 

118 return self.A * other 

119 

120 def __rmul__(self, other): 

121 return other * self.A 

122 

123 def __truediv__(self, other): 

124 return self.A / other 

125 

126 def __rtruediv__(self, other): 

127 return other / self.A 

128 

129 def __pow__(self, other): 

130 return self.A**other 

131 

132 def __matmul__(self, other): 

133 return self.A @ other 

134 

135 def __rmatmul__(self, other): 

136 return other @ self.A 

137 

138 def __iter__(self): 

139 return self.A.__iter__() 

140 

141 def __len__(self): 

142 return self.A.__len__() 

143 

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

145 """Get the norm of the matrix.""" 

146 if self.flag == "numpy": 

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

148 

149 if self.flag == "scipy_sparse": 

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

151 

152 if self.flag == "cupy_full": 

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

154 

155 if self.flag == "cupy_sparse": 

156 return csparse.linalg.norm(self.A, ord=order) ** power 

157 raise ValueError("Unknown flag.") 

158 

159 # def normalize_rows(self, order=None, power=1): 

160 # """Normalize the rows of the matrix with the "norm" norm of power.""" 

161 # if self.flag == "numpy": 

162 # return self.A / (np.linalg.norm(self.A, axis=1, ord=order) ** power)[:, None] 

163 

164 # elif self.flag == "scipy_sparse": 

165 # return ( 

166 # sparse.diags_array(1 / (sparse.linalg.norm(self.A, axis=1, ord=order) ** power)) 

167 # @ self.A 

168 # ) 

169 

170 # elif self.flag == "cupy_full": 

171 # return self.A / (cp.linalg.norm(self.A, axis=1, order=order) ** power)[:, None] 

172 

173 # elif self.flag == "cupy_sparse": 

174 # return ( 

175 # csparse.diags( 

176 # (1 / (csparse.linalg.norm(self.A, axis=1, ord=order) ** power)).ravel() 

177 # ) 

178 # @ self.A 

179 # ) 

180 

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

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

183 if self.flag == "numpy": 

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

185 

186 if self.flag == "scipy_sparse": 

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

188 

189 if self.flag == "cupy_full": 

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

191 

192 if self.flag == "cupy_sparse": 

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

194 

195 raise ValueError("Unknown flag.") 

196 

197 def single_map(self, x, i): 

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

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

200 return self.A[i] @ x 

201 

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

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

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

205 raise ValueError("Unknown flag.") 

206 

207 def index_map(self, x, idx): 

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

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

210 return self.A[idx] @ x 

211 

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

213 return self.A[idx] @ x 

214 

215 raise ValueError("Unknown flag.") 

216 

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

218 

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

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

221 

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

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

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

225 else: 

226 raise ValueError("Unknown flag.") 

227 

228 def getrow(self, i): 

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

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

231 return self.A[i] 

232 

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

234 return self.A.getrow(i) 

235 

236 raise ValueError("Unknown flag.")