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

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

2 import numpy as np

3 from scipy import sparse

4

5 try:

6 import cupy as cp

7 import cupyx.scipy.sparse as csparse

8

9 NO_GPU = False

10

11 except ImportError:

12 NO_GPU = True

13 cp = None

14 csparse = None

15

16

17 class 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.")