1# Copyright 2011-2013 Kwant authors. 

2# 

3# This file is part of Kwant. It is subject to the license terms in the file 

4# LICENSE.rst found in the top-level directory of this distribution and at 

5# https://kwant-project.org/license. A list of Kwant authors can be found in 

6# the file AUTHORS.rst at the top-level directory of this distribution and at 

7# https://kwant-project.org/authors. 

8 

9"""Interface to the MUMPS sparse solver library""" 

10 

11__all__ = ['MUMPSContext', 'schur_complement', 'AnalysisStatistics', 

12 'FactorizationStatistics', 'MUMPSError'] 

13 

14import time 

15import numpy as np 

16import scipy.sparse 

17import warnings 

18from . import _mumps 

19from .fortran_helpers import prepare_for_fortran 

20 

21orderings = { 'amd' : 0, 'amf' : 2, 'scotch' : 3, 'pord' : 4, 'metis' : 5, 

22 'qamd' : 6, 'auto' : 7 } 

23 

24ordering_name = [ 'amd', 'user-defined', 'amf', 

25 'scotch', 'pord', 'metis', 'qamd'] 

26 

27 

28def possible_orderings(): 

29 """Return the ordering options that are available in the current 

30 installation of MUMPS. 

31 

32 Which ordering options are actually available depends how MUMPs was 

33 compiled. Note that passing an ordering that is not avaialble in the 

34 current installation of MUMPS will not fail, instead MUMPS will fall back 

35 to a supported one. 

36 

37 Returns 

38 ------- 

39 orderings : list of strings 

40 A list of installed orderings that can be used in the `ordering` option 

41 of MUMPS. 

42 """ 

43 

44 if not possible_orderings.cached: 

45 # Try all orderings on a small test matrix, and check which one was 

46 # actually used. 

47 

48 possible_orderings.cached = ['auto'] 

49 for ordering in [0, 2, 3, 4, 5, 6]: 

50 data = np.asfortranarray([1, 1], dtype=np.complex128) 

51 row = np.asfortranarray([1, 2], dtype=_mumps.int_dtype) 

52 col = np.asfortranarray([1, 2], dtype=_mumps.int_dtype) 

53 

54 instance = _mumps.zmumps() 

55 instance.set_assembled_matrix(2, row, col, data) 

56 instance.icntl[7] = ordering 

57 instance.job = 1 

58 instance.call() 

59 

60 if instance.infog[7] == ordering: 60 ↛ 49line 60 didn't jump to line 49, because the condition on line 60 was never false

61 possible_orderings.cached.append(ordering_name[ordering]) 

62 

63 return possible_orderings.cached 

64 

65possible_orderings.cached = None 

66 

67 

68error_messages = { 

69 -5 : "Not enough memory during analysis phase", 

70 -6 : "Matrix is singular in structure", 

71 -7 : "Not enough memory during analysis phase", 

72 -10 : "Matrix is numerically singular", 

73 -11 : "The authors of MUMPS would like to hear about this", 

74 -12 : "The authors of MUMPS would like to hear about this", 

75 -13 : "Not enough memory" 

76} 

77 

78class MUMPSError(RuntimeError): 

79 def __init__(self, infog): 

80 self.error = infog[1] 

81 if self.error in error_messages: 

82 msg = "{}. (MUMPS error {})".format( 

83 error_messages[self.error], self.error) 

84 else: 

85 msg = "MUMPS failed with error {}.".format(self.error) 

86 

87 RuntimeError.__init__(self, msg) 

88 

89 

90class AnalysisStatistics: 

91 def __init__(self, inst, time=None): 

92 self.est_mem_incore = inst.infog[17] 

93 self.est_mem_ooc = inst.infog[27] 

94 self.est_nonzeros = (inst.infog[20] if inst.infog[20] > 0 else 

95 -inst.infog[20] * 1000000) 

96 self.est_flops = inst.rinfog[1] 

97 self.ordering = ordering_name[inst.infog[7]] 

98 self.time = time 

99 

100 def __str__(self): 

101 parts = ["estimated memory for in-core factorization:", 

102 str(self.est_mem_incore), "mbytes\n", 

103 "estimated memory for out-of-core factorization:", 

104 str(self.est_mem_ooc), "mbytes\n", 

105 "estimated number of nonzeros in factors:", 

106 str(self.est_nonzeros), "\n", 

107 "estimated number of flops:", str(self.est_flops), "\n", 

108 "ordering used:", self.ordering] 

109 if hasattr(self, "time"): 

110 parts.extend(["\n analysis time:", str(self.time), "secs"]) 

111 return " ".join(parts) 

112 

113 

114class FactorizationStatistics: 

115 def __init__(self, inst, time=None, include_ordering=False): 

116 # information about pivoting 

117 self.offdiag_pivots = inst.infog[12] if inst.sym == 0 else 0 

118 self.delayed_pivots = inst.infog[13] 

119 self.tiny_pivots = inst.infog[25] 

120 

121 # possibly include ordering (used in schur_complement) 

122 if include_ordering: 122 ↛ 123line 122 didn't jump to line 123, because the condition on line 122 was never true

123 self.ordering = ordering_name[inst.infog[7]] 

124 

125 # information about runtime effiency 

126 self.memory = inst.infog[22] 

127 self.nonzeros = (inst.infog[29] if inst.infog[29] > 0 else 

128 -inst.infog[29] * 1000000) 

129 self.flops = inst.rinfog[3] 

130 if time: 130 ↛ exitline 130 didn't return from function '__init__', because the condition on line 130 was never false

131 self.time = time 

132 

133 def __str__(self): 

134 parts = ["off-diagonal pivots:", str(self.offdiag_pivots), "\n", 

135 "delayed pivots:", str(self.delayed_pivots), "\n", 

136 "tiny pivots:", str(self.tiny_pivots), "\n"] 

137 if hasattr(self, "ordering"): 

138 parts.extend(["ordering used:", self.ordering, "\n"]) 

139 parts.extend(["memory used during factorization:", str(self.memory), 

140 "mbytes\n", 

141 "nonzeros in factored matrix:", str(self.nonzeros), "\n", 

142 "floating point operations:", str(self.flops)]) 

143 if hasattr(self, "time"): 

144 parts.extend(["\n factorization time:", str(self.time), "secs"]) 

145 return " ".join(parts) 

146 

147 

148class MUMPSContext: 

149 """MUMPSContext contains the internal data structures needed by the 

150 MUMPS library and contains a user-friendly interface. 

151 

152 WARNING: Only complex numbers supported. 

153 

154 Examples 

155 -------- 

156 

157 Solving a small system of equations. 

158 

159 >>> import scipy.sparse as sp 

160 >>> a = sp.coo_matrix([[1.,0],[0,2.]], dtype=complex) 

161 >>> ctx = kwant.linalg.mumps.MUMPSContext() 

162 >>> ctx.factor(a) 

163 >>> ctx.solve([1., 1.]) 

164 array([ 1.0+0.j, 0.5+0.j]) 

165 

166 Instance variables 

167 ------------------ 

168 

169 analysis_stats : `AnalysisStatistics` 

170 contains MUMPS statistics after an analysis step (i.e. after a call to 

171 `analyze` or `factor`) 

172 factor_stats : `FactorizationStatistics` 

173 contains MUMPS statistics after a factorization step (i.e. after a 

174 call to `factor`) 

175 

176 """ 

177 

178 def __init__(self, verbose=False): 

179 """Init the MUMPSContext class 

180 

181 Parameters 

182 ---------- 

183 

184 verbose : True or False 

185 control whether MUMPS prints lots of internal statistics 

186 and debug information to screen. 

187 """ 

188 self.mumps_instance = None 

189 self.dtype = None 

190 self.verbose = verbose 

191 self.factored = False 

192 

193 def analyze(self, a, ordering='auto', overwrite_a=False): 

194 """Perform analysis step of MUMPS. 

195 

196 In the analyis step, MUMPS figures out a reordering for the matrix and 

197 estimates number of operations and memory needed for the factorization 

198 time. This step usually needs not be called separately (it is done 

199 automatically by `factor`), but it can be useful to test which ordering 

200 would give best performance in the actual factorization, as MUMPS 

201 estimates are available in `analysis_stats`. 

202 

203 Parameters 

204 ---------- 

205 

206 a : sparse SciPy matrix 

207 input matrix. Internally, the matrix is converted to `coo` format 

208 (so passing this format is best for performance) 

209 ordering : { 'auto', 'amd', 'amf', 'scotch', 'pord', 'metis', 'qamd' } 

210 ordering to use in the factorization. The availability of a 

211 particular ordering depends on the MUMPS installation. Default is 

212 'auto'. 

213 overwrite_a : True or False 

214 whether the data in a may be overwritten, which can lead to a small 

215 performance gain. Default is False. 

216 """ 

217 

218 a = a.tocoo() 

219 

220 if a.ndim != 2 or a.shape[0] != a.shape[1]: 220 ↛ 221line 220 didn't jump to line 221, because the condition on line 220 was never true

221 raise ValueError("Input matrix must be square!") 

222 

223 if not ordering in orderings.keys(): 223 ↛ 224line 223 didn't jump to line 224, because the condition on line 223 was never true

224 raise ValueError("Unknown ordering '"+ordering+"'!") 

225 

226 dtype, row, col, data = _make_assembled_from_coo(a, overwrite_a) 

227 

228 if dtype != self.dtype: 228 ↛ 232line 228 didn't jump to line 232, because the condition on line 228 was never false

229 self.mumps_instance = getattr(_mumps, dtype+"mumps")(self.verbose) 

230 self.dtype = dtype 

231 

232 self.n = a.shape[0] 

233 self.row = row 

234 self.col = col 

235 self.data = data 

236 # Note: if I don't store them, they go out of scope and are 

237 # deleted. I however need the memory to stay around! 

238 

239 self.mumps_instance.set_assembled_matrix(a.shape[0], row, col, data) 

240 self.mumps_instance.icntl[7] = orderings[ordering] 

241 self.mumps_instance.job = 1 

242 t1 = time.process_time() 

243 self.mumps_instance.call() 

244 t2 = time.process_time() 

245 self.factored = False 

246 

247 if self.mumps_instance.infog[1] < 0: 247 ↛ 248line 247 didn't jump to line 248, because the condition on line 247 was never true

248 raise MUMPSError(self.mumps_instance.infog) 

249 

250 self.analysis_stats = AnalysisStatistics(self.mumps_instance, 

251 t2 - t1) 

252 

253 def factor(self, a, ordering='auto', ooc=False, pivot_tol=0.01, 

254 reuse_analysis=False, overwrite_a=False): 

255 """Perform the LU factorization of the matrix. 

256 

257 This LU factorization can then later be used to solve a linear system 

258 with `solve`. Statistical data of the factorization is stored in 

259 `factor_stats`. 

260 

261 Parameters 

262 ---------- 

263 

264 a : sparse SciPy matrix 

265 input matrix. Internally, the matrix is converted to `coo` format 

266 (so passing this format is best for performance) 

267 ordering : { 'auto', 'amd', 'amf', 'scotch', 'pord', 'metis', 'qamd' } 

268 ordering to use in the factorization. The availability of a 

269 particular ordering depends on the MUMPS installation. Default is 

270 'auto'. 

271 ooc : True or False 

272 whether to use the out-of-core functionality of MUMPS. 

273 (out-of-core means that data is written to disk to reduce memory 

274 usage.) Default is False. 

275 pivot_tol: number in the range [0, 1] 

276 pivoting threshold. Pivoting is typically limited in sparse 

277 solvers, as too much pivoting destroys sparsity. 1.0 means full 

278 pivoting, whereas 0.0 means no pivoting. Default is 0.01. 

279 reuse_analysis: True or False 

280 whether to reuse the analysis done in a previous call to `analyze` 

281 or `factor`. If the structure of the matrix stays the same, and the 

282 numerical values do not change much, the previous analysis can be 

283 reused, saving some time. WARNING: There is no check whether the 

284 structure of your matrix is compatible with the previous 

285 analysis. Also, if the values are not similar enough, there might 

286 be loss of accuracy, without a warning. Default is False. 

287 overwrite_a : True or False 

288 whether the data in a may be overwritten, which can lead to a small 

289 performance gain. Default is False. 

290 """ 

291 a = a.tocoo() 

292 

293 if a.ndim != 2 or a.shape[0] != a.shape[1]: 293 ↛ 294line 293 didn't jump to line 294, because the condition on line 293 was never true

294 raise ValueError("Input matrix must be square!") 

295 

296 # Analysis phase must be done before factorization 

297 # Note: previous analysis is reused only if reuse_analysis == True 

298 

299 if reuse_analysis: 

300 if self.mumps_instance is None: 300 ↛ 307line 300 didn't jump to line 307, because the condition on line 300 was never false

301 warnings.warn("Missing analysis although reuse_analysis=True. " 

302 "New analysis is performed.", 

303 RuntimeWarning, 

304 stacklevel=2) 

305 self.analyze(a, ordering=ordering, overwrite_a=overwrite_a) 

306 else: 

307 dtype, row, col, data = _make_assembled_from_coo(a, 

308 overwrite_a) 

309 if self.dtype != dtype: 

310 raise ValueError("MUMPSContext dtype and matrix dtype " 

311 "incompatible!") 

312 

313 self.n = a.shape[0] 

314 self.row = row 

315 self.col = col 

316 self.data = data 

317 self.mumps_instance.set_assembled_matrix(a.shape[0], 

318 row, col, data) 

319 else: 

320 self.analyze(a, ordering=ordering, overwrite_a=overwrite_a) 

321 

322 self.mumps_instance.icntl[22] = 1 if ooc else 0 

323 self.mumps_instance.job = 2 

324 self.mumps_instance.cntl[1] = pivot_tol 

325 

326 done = False 

327 while not done: 

328 t1 = time.process_time() 

329 self.mumps_instance.call() 

330 t2 = time.process_time() 

331 

332 # error -8, -9 (not enough allocated memory) is treated 

333 # specially, by increasing the memory relaxation parameter 

334 if self.mumps_instance.infog[1] < 0: 

335 if self.mumps_instance.infog[1] in (-8, -9): 335 ↛ 339line 335 didn't jump to line 339, because the condition on line 335 was never false

336 # double the additional memory 

337 self.mumps_instance.icntl[14] *= 2 

338 else: 

339 raise MUMPSError(self.mumps_instance.infog) 

340 else: 

341 done = True 

342 

343 self.factored = True 

344 self.factor_stats = FactorizationStatistics(self.mumps_instance, 

345 t2 - t1) 

346 

347 def _solve_sparse(self, b): 

348 b = b.tocsc() 

349 x = np.empty((b.shape[0], b.shape[1]), 

350 order='F', dtype=self.data.dtype) 

351 

352 dtype, col_ptr, row_ind, data = _make_sparse_rhs_from_csc( 

353 b, self.data.dtype) 

354 

355 if b.shape[0] != self.n: 355 ↛ 356line 355 didn't jump to line 356, because the condition on line 355 was never true

356 raise ValueError("Right hand side has wrong size") 

357 

358 if self.dtype != dtype: 358 ↛ 359line 358 didn't jump to line 359, because the condition on line 358 was never true

359 raise ValueError("Data type of right hand side is not " 

360 "compatible with the dtype of the " 

361 "linear system") 

362 

363 self.mumps_instance.set_sparse_rhs(col_ptr, row_ind, data) 

364 self.mumps_instance.set_dense_rhs(x) 

365 self.mumps_instance.job = 3 

366 self.mumps_instance.icntl[20] = 1 

367 self.mumps_instance.call() 

368 

369 return x 

370 

371 def _solve_dense(self, b, overwrite_b=False): 

372 dtype, b = prepare_for_fortran(overwrite_b, b, 

373 np.zeros(1, dtype=self.data.dtype))[:2] 

374 

375 if b.shape[0] != self.n: 375 ↛ 376line 375 didn't jump to line 376, because the condition on line 375 was never true

376 raise ValueError("Right hand side has wrong size") 

377 

378 if self.dtype != dtype: 378 ↛ 379line 378 didn't jump to line 379, because the condition on line 378 was never true

379 raise ValueError("Data type of right hand side is not " 

380 "compatible with the dtype of the " 

381 "linear system") 

382 

383 self.mumps_instance.set_dense_rhs(b) 

384 self.mumps_instance.job = 3 

385 self.mumps_instance.call() 

386 

387 return b 

388 

389 def solve(self, b, overwrite_b=False): 

390 """Solve a linear system after the LU factorization has previously 

391 been performed by `factor`. 

392 

393 Supports both dense and sparse right hand sides. 

394 

395 Parameters 

396 ---------- 

397 

398 b : dense (NumPy) matrix or vector or sparse (SciPy) matrix 

399 the right hand side to solve. Accepts both dense and sparse input; 

400 if the input is sparse 'csc' format is used internally (so passing 

401 a 'csc' matrix gives best performance). 

402 overwrite_b : True or False 

403 whether the data in b may be overwritten, which can lead to a small 

404 performance gain. Default is False. 

405 

406 Returns 

407 ------- 

408 

409 x : NumPy array 

410 the solution to the linear system as a dense matrix (a vector is 

411 returned if b was a vector, otherwise a matrix is returned). 

412 """ 

413 

414 if not self.factored: 414 ↛ 415line 414 didn't jump to line 415, because the condition on line 414 was never true

415 raise RuntimeError("Factorization must be done before solving!") 

416 

417 if scipy.sparse.isspmatrix(b): 

418 return self._solve_sparse(b) 

419 else: 

420 return self._solve_dense(b, overwrite_b) 

421 

422 

423def schur_complement(a, indices, ordering='auto', ooc=False, pivot_tol=0.01, 

424 calc_stats=False, overwrite_a=False): 

425 """Compute the Schur complement block of matrix a using MUMPS. 

426 

427 Parameters: 

428 a : sparse matrix 

429 input matrix. Internally, the matrix is converted to `coo` format (so 

430 passing this format is best for performance) 

431 indices : 1d array 

432 indices (row and column) of the desired Schur complement block. (The 

433 Schur complement block is square, so that the indices are both row and 

434 column indices.) 

435 ordering : { 'auto', 'amd', 'amf', 'scotch', 'pord', 'metis', 'qamd' } 

436 ordering to use in the factorization. The availability of a particular 

437 ordering depends on the MUMPS installation. Default is 'auto'. 

438 ooc : True or False 

439 whether to use the out-of-core functionality of MUMPS. (out-of-core 

440 means that data is written to disk to reduce memory usage.) Default is 

441 False. 

442 pivot_tol: number in the range [0, 1] 

443 pivoting threshold. Pivoting is typically limited in sparse solvers, as 

444 too much pivoting destroys sparsity. 1.0 means full pivoting, whereas 

445 0.0 means no pivoting. Default is 0.01. 

446 calc_stats: True or False 

447 whether to return the analysis and factorization statistics collected 

448 by MUMPS. Default is False. 

449 overwrite_a : True or False 

450 whether the data in a may be overwritten, which can lead to a small 

451 performance gain. Default is False. 

452 

453 Returns 

454 ------- 

455 

456 s : NumPy array 

457 Schur complement block 

458 factor_stats: `FactorizationStatistics` 

459 statistics of the factorization as collected by MUMPS. Only returned 

460 if ``calc_stats==True``. 

461 """ 

462 

463 if not scipy.sparse.isspmatrix(a): 463 ↛ 464line 463 didn't jump to line 464, because the condition on line 463 was never true

464 raise ValueError("a must be a sparse SciPy matrix!") 

465 

466 a = a.tocoo() 

467 

468 if a.ndim != 2 or a.shape[0] != a.shape[1]: 468 ↛ 469line 468 didn't jump to line 469, because the condition on line 468 was never true

469 raise ValueError("Input matrix must be square!") 

470 

471 indices = np.asanyarray(indices) 

472 

473 if indices.ndim != 1: 473 ↛ 474line 473 didn't jump to line 474, because the condition on line 473 was never true

474 raise ValueError("Schur indices must be specified in a 1d array!") 

475 

476 if not ordering in orderings.keys(): 476 ↛ 477line 476 didn't jump to line 477, because the condition on line 476 was never true

477 raise ValueError("Unknown ordering '"+ordering+"'!") 

478 

479 dtype, row, col, data = _make_assembled_from_coo(a, overwrite_a) 

480 indices = _make_mumps_index_array(indices) 

481 

482 mumps_instance = getattr(_mumps, dtype+"mumps")() 

483 

484 mumps_instance.set_assembled_matrix(a.shape[0], row, col, data) 

485 mumps_instance.icntl[7] = orderings[ordering] 

486 mumps_instance.icntl[19] = 1 

487 mumps_instance.icntl[31] = 1 # discard factors, from 4.10.0 

488 # has no effect in earlier versions 

489 

490 schur_compl = np.empty((indices.size, indices.size), 

491 order='C', dtype=data.dtype) 

492 mumps_instance.set_schur(schur_compl, indices) 

493 

494 mumps_instance.job = 4 # job=4 -> 1 and 2 after each other 

495 t1 = time.process_time() 

496 mumps_instance.call() 

497 t2 = time.process_time() 

498 

499 if not calc_stats: 499 ↛ 502line 499 didn't jump to line 502, because the condition on line 499 was never false

500 return schur_compl 

501 else: 

502 return (schur_compl, FactorizationStatistics( 

503 mumps_instance, time=t2 - t1, include_ordering=True)) 

504 

505 

506# Some internal helper functions 

507def _make_assembled_from_coo(a, overwrite_a): 

508 dtype, data = prepare_for_fortran(overwrite_a, a.data) 

509 

510 row = np.asfortranarray(a.row.astype(_mumps.int_dtype)) 

511 col = np.asfortranarray(a.col.astype(_mumps.int_dtype)) 

512 

513 # MUMPS uses Fortran indices. 

514 row += 1 

515 col += 1 

516 

517 return dtype, row, col, data 

518 

519 

520def _make_sparse_rhs_from_csc(b, dtype): 

521 dtype, data = prepare_for_fortran(True, b.data, 

522 np.zeros(1, dtype=dtype))[:2] 

523 

524 col_ptr = np.asfortranarray(b.indptr.astype(_mumps.int_dtype)) 

525 row_ind = np.asfortranarray(b.indices.astype(_mumps.int_dtype)) 

526 

527 # MUMPS uses Fortran indices. 

528 col_ptr += 1 

529 row_ind += 1 

530 

531 return dtype, col_ptr, row_ind, data 

532 

533 

534def _make_mumps_index_array(a): 

535 a = np.asfortranarray(a.astype(_mumps.int_dtype)) 

536 a += 1 # Fortran indices 

537 

538 return a