Coverage for src / sdynpy / core / sdynpy_matrix.py: 41%

1# -*- coding: utf-8 -*- 

2""" 

3Defines a matrix that has helpful tools for bookkeeping. 

4""" 

5""" 

6Copyright 2022 National Technology & Engineering Solutions of Sandia, 

7LLC (NTESS). Under the terms of Contract DE-NA0003525 with NTESS, the U.S. 

8Government retains certain rights in this software. 

9 

10This program is free software: you can redistribute it and/or modify 

11it under the terms of the GNU General Public License as published by 

12the Free Software Foundation, either version 3 of the License, or 

13(at your option) any later version. 

14 

15This program is distributed in the hope that it will be useful, 

16but WITHOUT ANY WARRANTY; without even the implied warranty of 

17MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

18GNU General Public License for more details. 

19 

20You should have received a copy of the GNU General Public License 

21along with this program. If not, see <https://www.gnu.org/licenses/>. 

22""" 

23 

24import numpy as np 

25import matplotlib.pyplot as plt 

26import matplotlib.ticker as ticker 

27from .sdynpy_array import SdynpyArray 

28from .sdynpy_coordinate import CoordinateArray, outer_product 

29 

30 

31class Matrix(SdynpyArray): 

32 """Matrix with degrees of freedom stored for better bookkeeping. 

33 

34 Use the matrix helper function to create the object. 

35 """ 

36 

37 @staticmethod 

38 def data_dtype(rows, columns, is_complex=False): 

39 """ 

40 Data type of the underlying numpy structured array for real shapes 

41 

42 Parameters 

43 ---------- 

44 rows : int 

45 Number of rows in the matrix 

46 columns : int 

47 Number of columns in the matrix 

48 

49 Returns 

50 ------- 

51 list 

52 Numpy dtype that can be passed into any of the numpy array 

53 constructors 

54 

55 """ 

56 return [ 

57 ('matrix', 'complex128' if is_complex else 'float64', (rows, columns)), 

58 ('row_coordinate', CoordinateArray.data_dtype, (rows,)), 

59 ('column_coordinate', CoordinateArray.data_dtype, (columns,)), 

60 ] 

61 

62 """Datatype for the underlying numpy structured array""" 

63 

64 def __new__(subtype, shape, nrows, ncols, is_complex=False, buffer=None, offset=0, 

65 strides=None, order=None): 

66 # Create the ndarray instance of our type, given the usual 

67 # ndarray input arguments. This will call the standard 

68 # ndarray constructor, but return an object of our type. 

69 # It also triggers a call to InfoArray.__array_finalize__ 

70 obj = super(Matrix, subtype).__new__(subtype, shape, 

71 Matrix.data_dtype(nrows, ncols, is_complex), 

72 buffer, offset, strides, 

73 order) 

74 # Finally, we must return the newly created object: 

75 return obj 

76 

77 @property 

78 def coordinate(self): 

79 """ 

80 Returns the full coordinate array for the matrix 

81 """ 

82 return outer_product(self.row_coordinate, self.column_coordinate) 

83 

84 @coordinate.setter 

85 def coordinate(self, value): 

86 raise RuntimeError('Cannot set coordinate directly. Set row_coordinate or column_coordinate instead.') 

87 

88 @property 

89 def num_coordinate_rows(self): 

90 """ 

91 Returns the number of coordinate rows 

92 """ 

93 return self.matrix.shape[-2] 

94 

95 @property 

96 def num_coordinate_columns(self): 

97 """ 

98 Returns the number of coordinate columns 

99 """ 

100 return self.matrix.shape[-1] 

101 

102 def __setitem__(self, key, value): 

103 if type(key) is tuple and len(key) == 2 and any([(type(val) is CoordinateArray) for val in key]): 

104 # Get indices corresponding to the coordinates 

105 row_request, column_request = key 

106 if row_request is None or (isinstance(row_request, slice) and row_request == slice(None)) or isinstance(row_request, type(Ellipsis)): 

107 matrix_row_indices = np.arange(self.num_coordinate_rows) 

108 request_row_indices = np.arange(self.num_coordinate_rows) 

109 row_multiplications = np.ones(self.num_coordinate_rows) 

110 else: 

111 row_request = np.atleast_1d(row_request) 

112 column_request = np.atleast_1d(column_request) 

113 # Start with rows 

114 coordinate_array = self.row_coordinate 

115 single_matrix_coordinate_array = coordinate_array[ 

116 (0,) * (coordinate_array.ndim - 1) + (slice(None),)] 

117 # Now check if the coordinates are consistent across the arrays 

118 if not np.all((coordinate_array[..., :] == single_matrix_coordinate_array)): 

119 # If they aren't, raise a value error 

120 raise ValueError( 

121 'Matrix must have equivalent row coordinates for all matrices to index by coordinate') 

122 consistent_row_coordinates, matrix_row_indices, request_row_indices = np.intersect1d( 

123 abs(single_matrix_coordinate_array), abs(row_request), assume_unique=False, return_indices=True) 

124 if consistent_row_coordinates.size != row_request.size: 

125 extra_keys = np.setdiff1d(abs(row_request), abs(single_matrix_coordinate_array)) 

126 if extra_keys.size == 0: 

127 raise ValueError( 

128 'Duplicate coordinate values requested. Please ensure coordinate indices are unique.') 

129 raise ValueError( 

130 'Not all indices in requested coordinate array exist in the shape\n{:}'.format(str(extra_keys))) 

131 # Handle sign flipping 

132 row_multiplications = row_request.flatten()[request_row_indices].sign( 

133 ) * single_matrix_coordinate_array[matrix_row_indices].sign() 

134 # # Invert the indices to return the dofs in the correct order as specified in keys 

135 # inverse_row_indices = np.zeros(matrix_row_indices.shape, dtype=int) 

136 # inverse_row_indices[matrix_row_indices] = np.arange(len(matrix_row_indices)) 

137 if column_request is None or (isinstance(column_request, slice) and column_request == slice(None)) or isinstance(column_request, type(Ellipsis)): 

138 matrix_col_indices = np.arange(self.num_coordinate_columns) 

139 request_col_indices = np.arange(self.num_coordinate_columns) 

140 col_multiplications = np.ones(self.num_coordinate_columns) 

141 else: 

142 # Now columns 

143 coordinate_array = self.column_coordinate 

144 single_matrix_coordinate_array = coordinate_array[ 

145 (0,) * (coordinate_array.ndim - 1) + (slice(None),)] 

146 # Now check if the coordinates are consistent across the arrays 

147 if not np.all((coordinate_array[..., :] == single_matrix_coordinate_array)): 

148 # If they aren't, raise a value error 

149 raise ValueError( 

150 'Matrix must have equivalent column coordinates for all matrices to index by coordinate') 

151 consistent_col_coordinates, matrix_col_indices, request_col_indices = np.intersect1d( 

152 abs(single_matrix_coordinate_array), abs(column_request), assume_unique=False, return_indices=True) 

153 if consistent_col_coordinates.size != column_request.size: 

154 extra_keys = np.setdiff1d(abs(column_request), abs(single_matrix_coordinate_array)) 

155 if extra_keys.size == 0: 

156 raise ValueError( 

157 'Duplicate coordinate values requested. Please ensure coordinate indices are unique.') 

158 raise ValueError( 

159 'Not all indices in requested coordinate array exist in the shape\n{:}'.format(str(extra_keys))) 

160 # Handle sign flipping 

161 col_multiplications = column_request.flatten()[request_col_indices].sign( 

162 ) * single_matrix_coordinate_array[matrix_col_indices].sign() 

163 # # Invert the indices to return the dofs in the correct order as specified in keys 

164 # inverse_col_indices = np.zeros(matrix_col_indices.shape, dtype=int) 

165 # inverse_col_indices[matrix_col_indices] = np.arange(len(matrix_col_indices)) 

166 value = np.array(value) 

167 value = np.broadcast_to(value, 

168 value.shape[:-2] + (len(consistent_row_coordinates), 

169 len(consistent_col_coordinates))) 

170 self.matrix[..., matrix_row_indices[:, np.newaxis], matrix_col_indices] = ( 

171 value[..., request_row_indices[:, np.newaxis], request_col_indices] 

172 * row_multiplications[:, np.newaxis] * col_multiplications) 

173 else: 

174 super().__setitem__(key, value) 

175 

176 def __getitem__(self, key): 

177 if type(key) is tuple and len(key) == 2 and any([(type(val) is CoordinateArray) for val in key]): 

178 # Get indices corresponding to the coordinates 

179 row_request, column_request = key 

180 if row_request is None or (isinstance(row_request, slice) and row_request == slice(None)) or isinstance(row_request, type(Ellipsis)): 

181 matrix_row_indices = np.arange(self.num_coordinate_rows) 

182 inverse_row_indices = np.arange(self.num_coordinate_rows) 

183 row_multiplications = np.ones(self.num_coordinate_rows) 

184 else: 

185 row_request = np.atleast_1d(row_request) 

186 column_request = np.atleast_1d(column_request) 

187 # Start with rows 

188 coordinate_array = self.row_coordinate 

189 single_matrix_coordinate_array = coordinate_array[ 

190 (0,) * (coordinate_array.ndim - 1) + (slice(None),)] 

191 # Now check if the coordinates are consistent across the arrays 

192 if not np.all((coordinate_array[..., :] == single_matrix_coordinate_array)): 

193 # If they aren't, raise a value error 

194 raise ValueError( 

195 'Matrix must have equivalent row coordinates for all matrices to index by coordinate') 

196 consistent_row_coordinates, matrix_row_indices, request_row_indices = np.intersect1d( 

197 abs(single_matrix_coordinate_array), abs(row_request), assume_unique=False, return_indices=True) 

198 if consistent_row_coordinates.size != row_request.size: 

199 extra_keys = np.setdiff1d(abs(row_request), abs(single_matrix_coordinate_array)) 

200 if extra_keys.size == 0: 

201 raise ValueError( 

202 'Duplicate coordinate values requested. Please ensure coordinate indices are unique.') 

203 raise ValueError( 

204 'Not all indices in requested coordinate array exist in the shape\n{:}'.format(str(extra_keys))) 

205 # Handle sign flipping 

206 row_multiplications = row_request.flatten()[request_row_indices].sign( 

207 ) * single_matrix_coordinate_array[matrix_row_indices].sign() 

208 # Invert the indices to return the dofs in the correct order as specified in keys 

209 inverse_row_indices = np.zeros(request_row_indices.shape, dtype=int) 

210 inverse_row_indices[request_row_indices] = np.arange(len(request_row_indices)) 

211 if column_request is None or (isinstance(column_request, slice) and column_request == slice(None)) or isinstance(column_request, type(Ellipsis)): 

212 matrix_col_indices = np.arange(self.num_coordinate_columns) 

213 inverse_col_indices = np.arange(self.num_coordinate_columns) 

214 col_multiplications = np.ones(self.num_coordinate_columns) 

215 else: 

216 # Now columns 

217 coordinate_array = self.column_coordinate 

218 single_matrix_coordinate_array = coordinate_array[ 

219 (0,) * (coordinate_array.ndim - 1) + (slice(None),)] 

220 # Now check if the coordinates are consistent across the arrays 

221 if not np.all((coordinate_array[..., :] == single_matrix_coordinate_array)): 

222 # If they aren't, raise a value error 

223 raise ValueError( 

224 'Matrix must have equivalent column coordinates for all matrices to index by coordinate') 

225 consistent_col_coordinates, matrix_col_indices, request_col_indices = np.intersect1d( 

226 abs(single_matrix_coordinate_array), abs(column_request), assume_unique=False, return_indices=True) 

227 if consistent_col_coordinates.size != column_request.size: 

228 extra_keys = np.setdiff1d(abs(column_request), abs(single_matrix_coordinate_array)) 

229 if extra_keys.size == 0: 

230 raise ValueError( 

231 'Duplicate coordinate values requested. Please ensure coordinate indices are unique.') 

232 raise ValueError( 

233 'Not all indices in requested coordinate array exist in the shape\n{:}'.format(str(extra_keys))) 

234 # Handle sign flipping 

235 col_multiplications = column_request.flatten()[request_col_indices].sign( 

236 ) * single_matrix_coordinate_array[matrix_col_indices].sign() 

237 # Invert the indices to return the dofs in the correct order as specified in keys 

238 inverse_col_indices = np.zeros(request_col_indices.shape, dtype=int) 

239 inverse_col_indices[request_col_indices] = np.arange(len(request_col_indices)) 

240 return_value = self.matrix[..., matrix_row_indices[:, np.newaxis], matrix_col_indices] * row_multiplications[:, np.newaxis] * col_multiplications 

241 return_value = return_value[..., inverse_row_indices[:, np.newaxis], inverse_col_indices] 

242 return return_value 

243 else: 

244 output = super().__getitem__(key) 

245 if isinstance(key, str) and key in ['row_coordinate', 'column_coordinate']: 

246 return output.view(CoordinateArray) 

247 else: 

248 return output 

249 

250 def __repr__(self): 

251 return '\n\n'.join(['matrix = \n'+repr(self.matrix), 

252 'row coordinates = \n'+repr(self.row_coordinate), 

253 'column coordinates = \n'+repr(self.column_coordinate)]) 

254 

255 def argsort_coordinate(self): 

256 """ 

257 Returns indices used to sort the coordinates on the rows and columns 

258 

259 Returns 

260 ------- 

261 row_indices 

262 Indices used to sort the row coordinates. 

263 column_indices 

264 Indices used to sort the column coordinates 

265 

266 """ 

267 row_indices = np.empty(self.row_coordinate.shape, dtype=int) 

268 column_indices = np.empty(self.column_coordinate.shape, dtype=int) 

269 for key, matrix in self.ndenumerate(): 

270 row_indices[key] = np.argsort(self.row_coordinate[key]) 

271 column_indices[key] = np.argsort(self.column_coordinate[key]) 

272 return row_indices[..., :, np.newaxis], column_indices[..., np.newaxis, :] 

273 

274 def sort_coordinate(self): 

275 """ 

276 Returns a copy of the Matrix with coordinate sorted 

277 

278 Returns 

279 ------- 

280 return_val : Matrix 

281 Matrix with row and column coordinates sorted. 

282 

283 """ 

284 sorting = self.argsort_coordinate() 

285 return_val = self.copy() 

286 for key, matrix in self.ndenumerate(): 

287 return_val.matrix[key] = self.matrix[key][..., sorting[0][key], sorting[1][key]] 

288 return_val.row_coordinate = self.row_coordinate[..., sorting[0][key][:, 0]] 

289 return_val.column_coordinate = self.column_coordinate[..., sorting[1][key][0, :]] 

290 return return_val 

291 

292 def plot(self, ax=None, show_colorbar=True, **imshow_args): 

293 """ 

294 Plots the matrix with coordinates labeled 

295 

296 Parameters 

297 ---------- 

298 ax : axis, optional 

299 Axis on which the matrix will be plotted. If not specified, a new 

300 figure will be created. 

301 show_colorbar : bool, optional 

302 If true, show a colorbar. Default is true. 

303 **imshow_args : various 

304 Additional arguments passed to imshow 

305 

306 Raises 

307 ------ 

308 ValueError 

309 Raised if matrix is multi-dimensional. 

310 

311 """ 

312 if self.size > 1: 

313 raise ValueError('Cannot plot more than one Matrix object simultaneously') 

314 if ax is None: 

315 fig, ax = plt.subplots() 

316 

317 @ticker.FuncFormatter 

318 def row_formatter(x, pos): 

319 if int(x) < 0 or int(x) >= len(self.row_coordinate): 

320 return '' 

321 try: 

322 return str(self.row_coordinate[int(x)]) 

323 except IndexError: 

324 return '' 

325 

326 @ticker.FuncFormatter 

327 def col_formatter(x, pos): 

328 if int(x) < 0 or int(x) >= len(self.column_coordinate): 

329 return '' 

330 try: 

331 return str(self.column_coordinate[int(x)]) 

332 except IndexError: 

333 return '' 

334 out = ax.imshow(self.matrix, **imshow_args) 

335 if show_colorbar: 

336 plt.colorbar(out, ax=ax) 

337 ax.xaxis.set_major_formatter(col_formatter) 

338 ax.xaxis.set_major_locator(ticker.MaxNLocator(integer=True)) 

339 ax.yaxis.set_major_formatter(row_formatter) 

340 ax.yaxis.set_major_locator(ticker.MaxNLocator(integer=True)) 

341 

342 @classmethod 

343 def eye(cls, row_coordinate, column_coordinate=None): 

344 """ 

345 Creates an identity matrix with the specified coordinates 

346 

347 Parameters 

348 ---------- 

349 row_coordinate : CoordinateArray 

350 Coordinate array to use as the row coordinates 

351 column_coordinate : CoordinateArray, optional 

352 Coordinate array to use as the column coordinates. If not specified, 

353 the row coordinates are used. 

354 

355 Returns 

356 ------- 

357 Matrix 

358 Diagonal matrix with the specified coordinates 

359 

360 """ 

361 if column_coordinate is None: 

362 column_coordinate = row_coordinate 

363 return matrix(np.eye(row_coordinate.shape[-1], column_coordinate.shape[-1]), 

364 row_coordinate, column_coordinate) 

365 

366 def pinv(self, **pinv_params): 

367 """ 

368 Creates the pseudoinverse of the matrix 

369 

370 Parameters 

371 ---------- 

372 **pinv_params : various 

373 Extra keyword arguments are passed directly to np.linalg.pinv 

374 

375 Returns 

376 ------- 

377 Matrix 

378 A matrix consisting of the pseudoinverse of the original matrix 

379 

380 """ 

381 mat = np.linalg.pinv(self.matrix,**pinv_params) 

382 return matrix(matrix=mat,row_coordinate = self.column_coordinate, 

383 column_coordinate = self.row_coordinate) 

384 

385 def __matmul__(self,other): 

386 if not isinstance(other,Matrix): 

387 # If it is not another matrix, rely on the other data to define what 

388 # matrix multiplication means 

389 return NotImplemented 

390 return super().__matmul__(other) 

391 

392 @property 

393 def T(self): 

394 return matrix(np.moveaxis(self.matrix, -1, -2), self.column_coordinate, self.row_coordinate) 

395 

396 @property 

397 def H(self): 

398 return matrix( 

399 np.moveaxis(self.matrix, -1, -2).conjugate(), 

400 self.column_coordinate, 

401 self.row_coordinate, 

402 ) 

403 

404 

405def matrix(matrix, row_coordinate, column_coordinate, buffer=None, offset=0, 

406 strides=None, order=None): 

407 """ 

408 Create a matrix object 

409 

410 Parameters 

411 ---------- 

412 matrix : ndarray 

413 The values in the matrix object 

414 row_coordinate : CoordinateArray 

415 Coordinates to assign to the rows of the matrix 

416 column_coordinate : CoordinateArray 

417 Coordinates to assign to the columns of the matrix 

418 

419 Returns 

420 ------- 

421 matrix : Matrix 

422 Matrix object 

423 

424 """ 

425 shape = matrix.shape[:-2] 

426 nrow, ncol = matrix.shape[-2:] 

427 m = Matrix(shape, nrow, ncol, np.iscomplexobj(matrix), 

428 buffer=buffer, offset=offset, 

429 strides=strides, order=order) 

430 m.row_coordinate = row_coordinate 

431 m.column_coordinate = column_coordinate 

432 m.matrix = matrix 

433 return m