Coverage for /usr/lib/python3/dist-packages/scipy/sparse/linalg/_dsolve/linsolve.py: 11%

225 statements  

« prev     ^ index     » next       coverage.py v7.9.1, created at 2025-06-14 15:55 +0200

1from warnings import warn 

2 

3import numpy as np 

4from numpy import asarray 

5from scipy.sparse import (issparse, 

6 SparseEfficiencyWarning, csc_matrix, csr_matrix) 

7from scipy.sparse._sputils import is_pydata_spmatrix 

8from scipy.linalg import LinAlgError 

9import copy 

10 

11from . import _superlu 

12 

13noScikit = False 

14try: 

15 import scikits.umfpack as umfpack 

16except ImportError: 

17 noScikit = True 

18 

19useUmfpack = not noScikit 

20 

21__all__ = ['use_solver', 'spsolve', 'splu', 'spilu', 'factorized', 

22 'MatrixRankWarning', 'spsolve_triangular'] 

23 

24 

25class MatrixRankWarning(UserWarning): 

26 pass 

27 

28 

29def use_solver(**kwargs): 

30 """ 

31 Select default sparse direct solver to be used. 

32 

33 Parameters 

34 ---------- 

35 useUmfpack : bool, optional 

36 Use UMFPACK [1]_, [2]_, [3]_, [4]_. over SuperLU. Has effect only 

37 if ``scikits.umfpack`` is installed. Default: True 

38 assumeSortedIndices : bool, optional 

39 Allow UMFPACK to skip the step of sorting indices for a CSR/CSC matrix. 

40 Has effect only if useUmfpack is True and ``scikits.umfpack`` is 

41 installed. Default: False 

42 

43 Notes 

44 ----- 

45 The default sparse solver is UMFPACK when available 

46 (``scikits.umfpack`` is installed). This can be changed by passing 

47 useUmfpack = False, which then causes the always present SuperLU 

48 based solver to be used. 

49 

50 UMFPACK requires a CSR/CSC matrix to have sorted column/row indices. If 

51 sure that the matrix fulfills this, pass ``assumeSortedIndices=True`` 

52 to gain some speed. 

53 

54 References 

55 ---------- 

56 .. [1] T. A. Davis, Algorithm 832: UMFPACK - an unsymmetric-pattern 

57 multifrontal method with a column pre-ordering strategy, ACM 

58 Trans. on Mathematical Software, 30(2), 2004, pp. 196--199. 

59 https://dl.acm.org/doi/abs/10.1145/992200.992206 

60 

61 .. [2] T. A. Davis, A column pre-ordering strategy for the 

62 unsymmetric-pattern multifrontal method, ACM Trans. 

63 on Mathematical Software, 30(2), 2004, pp. 165--195. 

64 https://dl.acm.org/doi/abs/10.1145/992200.992205 

65 

66 .. [3] T. A. Davis and I. S. Duff, A combined unifrontal/multifrontal 

67 method for unsymmetric sparse matrices, ACM Trans. on 

68 Mathematical Software, 25(1), 1999, pp. 1--19. 

69 https://doi.org/10.1145/305658.287640 

70 

71 .. [4] T. A. Davis and I. S. Duff, An unsymmetric-pattern multifrontal 

72 method for sparse LU factorization, SIAM J. Matrix Analysis and 

73 Computations, 18(1), 1997, pp. 140--158. 

74 https://doi.org/10.1137/S0895479894246905T. 

75 

76 Examples 

77 -------- 

78 >>> import numpy as np 

79 >>> from scipy.sparse.linalg import use_solver, spsolve 

80 >>> from scipy.sparse import csc_matrix 

81 >>> R = np.random.randn(5, 5) 

82 >>> A = csc_matrix(R) 

83 >>> b = np.random.randn(5) 

84 >>> use_solver(useUmfpack=False) # enforce superLU over UMFPACK 

85 >>> x = spsolve(A, b) 

86 >>> np.allclose(A.dot(x), b) 

87 True 

88 >>> use_solver(useUmfpack=True) # reset umfPack usage to default 

89 """ 

90 if 'useUmfpack' in kwargs: 

91 globals()['useUmfpack'] = kwargs['useUmfpack'] 

92 if useUmfpack and 'assumeSortedIndices' in kwargs: 

93 umfpack.configure(assumeSortedIndices=kwargs['assumeSortedIndices']) 

94 

95def _get_umf_family(A): 

96 """Get umfpack family string given the sparse matrix dtype.""" 

97 _families = { 

98 (np.float64, np.int32): 'di', 

99 (np.complex128, np.int32): 'zi', 

100 (np.float64, np.int64): 'dl', 

101 (np.complex128, np.int64): 'zl' 

102 } 

103 

104 f_type = np.sctypeDict[A.dtype.name] 

105 i_type = np.sctypeDict[A.indices.dtype.name] 

106 

107 try: 

108 family = _families[(f_type, i_type)] 

109 

110 except KeyError as e: 

111 msg = 'only float64 or complex128 matrices with int32 or int64' \ 

112 ' indices are supported! (got: matrix: %s, indices: %s)' \ 

113 % (f_type, i_type) 

114 raise ValueError(msg) from e 

115 

116 # See gh-8278. Considered converting only if 

117 # A.shape[0]*A.shape[1] > np.iinfo(np.int32).max, 

118 # but that didn't always fix the issue. 

119 family = family[0] + "l" 

120 A_new = copy.copy(A) 

121 A_new.indptr = np.array(A.indptr, copy=False, dtype=np.int64) 

122 A_new.indices = np.array(A.indices, copy=False, dtype=np.int64) 

123 

124 return family, A_new 

125 

126def _safe_downcast_indices(A): 

127 # check for safe downcasting 

128 max_value = np.iinfo(np.intc).max 

129 

130 if A.indptr[-1] > max_value: # indptr[-1] is max b/c indptr always sorted 

131 raise ValueError("indptr values too large for SuperLU") 

132 

133 if max(*A.shape) > max_value: # only check large enough arrays 

134 if np.any(A.indices > max_value): 

135 raise ValueError("indices values too large for SuperLU") 

136 

137 indices = A.indices.astype(np.intc, copy=False) 

138 indptr = A.indptr.astype(np.intc, copy=False) 

139 return indices, indptr 

140 

141def spsolve(A, b, permc_spec=None, use_umfpack=True): 

142 """Solve the sparse linear system Ax=b, where b may be a vector or a matrix. 

143 

144 Parameters 

145 ---------- 

146 A : ndarray or sparse matrix 

147 The square matrix A will be converted into CSC or CSR form 

148 b : ndarray or sparse matrix 

149 The matrix or vector representing the right hand side of the equation. 

150 If a vector, b.shape must be (n,) or (n, 1). 

151 permc_spec : str, optional 

152 How to permute the columns of the matrix for sparsity preservation. 

153 (default: 'COLAMD') 

154 

155 - ``NATURAL``: natural ordering. 

156 - ``MMD_ATA``: minimum degree ordering on the structure of A^T A. 

157 - ``MMD_AT_PLUS_A``: minimum degree ordering on the structure of A^T+A. 

158 - ``COLAMD``: approximate minimum degree column ordering [1]_, [2]_. 

159 

160 use_umfpack : bool, optional 

161 if True (default) then use UMFPACK for the solution [3]_, [4]_, [5]_, 

162 [6]_ . This is only referenced if b is a vector and 

163 ``scikits.umfpack`` is installed. 

164 

165 Returns 

166 ------- 

167 x : ndarray or sparse matrix 

168 the solution of the sparse linear equation. 

169 If b is a vector, then x is a vector of size A.shape[1] 

170 If b is a matrix, then x is a matrix of size (A.shape[1], b.shape[1]) 

171 

172 Notes 

173 ----- 

174 For solving the matrix expression AX = B, this solver assumes the resulting 

175 matrix X is sparse, as is often the case for very sparse inputs. If the 

176 resulting X is dense, the construction of this sparse result will be 

177 relatively expensive. In that case, consider converting A to a dense 

178 matrix and using scipy.linalg.solve or its variants. 

179 

180 References 

181 ---------- 

182 .. [1] T. A. Davis, J. R. Gilbert, S. Larimore, E. Ng, Algorithm 836: 

183 COLAMD, an approximate column minimum degree ordering algorithm, 

184 ACM Trans. on Mathematical Software, 30(3), 2004, pp. 377--380. 

185 :doi:`10.1145/1024074.1024080` 

186 

187 .. [2] T. A. Davis, J. R. Gilbert, S. Larimore, E. Ng, A column approximate 

188 minimum degree ordering algorithm, ACM Trans. on Mathematical 

189 Software, 30(3), 2004, pp. 353--376. :doi:`10.1145/1024074.1024079` 

190 

191 .. [3] T. A. Davis, Algorithm 832: UMFPACK - an unsymmetric-pattern 

192 multifrontal method with a column pre-ordering strategy, ACM 

193 Trans. on Mathematical Software, 30(2), 2004, pp. 196--199. 

194 https://dl.acm.org/doi/abs/10.1145/992200.992206 

195 

196 .. [4] T. A. Davis, A column pre-ordering strategy for the 

197 unsymmetric-pattern multifrontal method, ACM Trans. 

198 on Mathematical Software, 30(2), 2004, pp. 165--195. 

199 https://dl.acm.org/doi/abs/10.1145/992200.992205 

200 

201 .. [5] T. A. Davis and I. S. Duff, A combined unifrontal/multifrontal 

202 method for unsymmetric sparse matrices, ACM Trans. on 

203 Mathematical Software, 25(1), 1999, pp. 1--19. 

204 https://doi.org/10.1145/305658.287640 

205 

206 .. [6] T. A. Davis and I. S. Duff, An unsymmetric-pattern multifrontal 

207 method for sparse LU factorization, SIAM J. Matrix Analysis and 

208 Computations, 18(1), 1997, pp. 140--158. 

209 https://doi.org/10.1137/S0895479894246905T. 

210 

211 

212 Examples 

213 -------- 

214 >>> import numpy as np 

215 >>> from scipy.sparse import csc_matrix 

216 >>> from scipy.sparse.linalg import spsolve 

217 >>> A = csc_matrix([[3, 2, 0], [1, -1, 0], [0, 5, 1]], dtype=float) 

218 >>> B = csc_matrix([[2, 0], [-1, 0], [2, 0]], dtype=float) 

219 >>> x = spsolve(A, B) 

220 >>> np.allclose(A.dot(x).toarray(), B.toarray()) 

221 True 

222 """ 

223 

224 if is_pydata_spmatrix(A): 

225 A = A.to_scipy_sparse().tocsc() 

226 

227 if not (issparse(A) and A.format in ("csc", "csr")): 

228 A = csc_matrix(A) 

229 warn('spsolve requires A be CSC or CSR matrix format', 

230 SparseEfficiencyWarning) 

231 

232 # b is a vector only if b have shape (n,) or (n, 1) 

233 b_is_sparse = issparse(b) or is_pydata_spmatrix(b) 

234 if not b_is_sparse: 

235 b = asarray(b) 

236 b_is_vector = ((b.ndim == 1) or (b.ndim == 2 and b.shape[1] == 1)) 

237 

238 # sum duplicates for non-canonical format 

239 A.sum_duplicates() 

240 A = A._asfptype() # upcast to a floating point format 

241 result_dtype = np.promote_types(A.dtype, b.dtype) 

242 if A.dtype != result_dtype: 

243 A = A.astype(result_dtype) 

244 if b.dtype != result_dtype: 

245 b = b.astype(result_dtype) 

246 

247 # validate input shapes 

248 M, N = A.shape 

249 if (M != N): 

250 raise ValueError(f"matrix must be square (has shape {(M, N)})") 

251 

252 if M != b.shape[0]: 

253 raise ValueError("matrix - rhs dimension mismatch (%s - %s)" 

254 % (A.shape, b.shape[0])) 

255 

256 use_umfpack = use_umfpack and useUmfpack 

257 

258 if b_is_vector and use_umfpack: 

259 if b_is_sparse: 

260 b_vec = b.toarray() 

261 else: 

262 b_vec = b 

263 b_vec = asarray(b_vec, dtype=A.dtype).ravel() 

264 

265 if noScikit: 

266 raise RuntimeError('Scikits.umfpack not installed.') 

267 

268 if A.dtype.char not in 'dD': 

269 raise ValueError("convert matrix data to double, please, using" 

270 " .astype(), or set linsolve.useUmfpack = False") 

271 

272 umf_family, A = _get_umf_family(A) 

273 umf = umfpack.UmfpackContext(umf_family) 

274 x = umf.linsolve(umfpack.UMFPACK_A, A, b_vec, 

275 autoTranspose=True) 

276 else: 

277 if b_is_vector and b_is_sparse: 

278 b = b.toarray() 

279 b_is_sparse = False 

280 

281 if not b_is_sparse: 

282 if A.format == "csc": 

283 flag = 1 # CSC format 

284 else: 

285 flag = 0 # CSR format 

286 

287 indices = A.indices.astype(np.intc, copy=False) 

288 indptr = A.indptr.astype(np.intc, copy=False) 

289 options = dict(ColPerm=permc_spec) 

290 x, info = _superlu.gssv(N, A.nnz, A.data, indices, indptr, 

291 b, flag, options=options) 

292 if info != 0: 

293 warn("Matrix is exactly singular", MatrixRankWarning) 

294 x.fill(np.nan) 

295 if b_is_vector: 

296 x = x.ravel() 

297 else: 

298 # b is sparse 

299 Afactsolve = factorized(A) 

300 

301 if not (b.format == "csc" or is_pydata_spmatrix(b)): 

302 warn('spsolve is more efficient when sparse b ' 

303 'is in the CSC matrix format', SparseEfficiencyWarning) 

304 b = csc_matrix(b) 

305 

306 # Create a sparse output matrix by repeatedly applying 

307 # the sparse factorization to solve columns of b. 

308 data_segs = [] 

309 row_segs = [] 

310 col_segs = [] 

311 for j in range(b.shape[1]): 

312 # TODO: replace this with 

313 # bj = b[:, j].toarray().ravel() 

314 # once 1D sparse arrays are supported. 

315 # That is a slightly faster code path. 

316 bj = b[:, [j]].toarray().ravel() 

317 xj = Afactsolve(bj) 

318 w = np.flatnonzero(xj) 

319 segment_length = w.shape[0] 

320 row_segs.append(w) 

321 col_segs.append(np.full(segment_length, j, dtype=int)) 

322 data_segs.append(np.asarray(xj[w], dtype=A.dtype)) 

323 sparse_data = np.concatenate(data_segs) 

324 sparse_row = np.concatenate(row_segs) 

325 sparse_col = np.concatenate(col_segs) 

326 x = A.__class__((sparse_data, (sparse_row, sparse_col)), 

327 shape=b.shape, dtype=A.dtype) 

328 

329 if is_pydata_spmatrix(b): 

330 x = b.__class__(x) 

331 

332 return x 

333 

334 

335def splu(A, permc_spec=None, diag_pivot_thresh=None, 

336 relax=None, panel_size=None, options=dict()): 

337 """ 

338 Compute the LU decomposition of a sparse, square matrix. 

339 

340 Parameters 

341 ---------- 

342 A : sparse matrix 

343 Sparse matrix to factorize. Most efficient when provided in CSC 

344 format. Other formats will be converted to CSC before factorization. 

345 permc_spec : str, optional 

346 How to permute the columns of the matrix for sparsity preservation. 

347 (default: 'COLAMD') 

348 

349 - ``NATURAL``: natural ordering. 

350 - ``MMD_ATA``: minimum degree ordering on the structure of A^T A. 

351 - ``MMD_AT_PLUS_A``: minimum degree ordering on the structure of A^T+A. 

352 - ``COLAMD``: approximate minimum degree column ordering 

353 

354 diag_pivot_thresh : float, optional 

355 Threshold used for a diagonal entry to be an acceptable pivot. 

356 See SuperLU user's guide for details [1]_ 

357 relax : int, optional 

358 Expert option for customizing the degree of relaxing supernodes. 

359 See SuperLU user's guide for details [1]_ 

360 panel_size : int, optional 

361 Expert option for customizing the panel size. 

362 See SuperLU user's guide for details [1]_ 

363 options : dict, optional 

364 Dictionary containing additional expert options to SuperLU. 

365 See SuperLU user guide [1]_ (section 2.4 on the 'Options' argument) 

366 for more details. For example, you can specify 

367 ``options=dict(Equil=False, IterRefine='SINGLE'))`` 

368 to turn equilibration off and perform a single iterative refinement. 

369 

370 Returns 

371 ------- 

372 invA : scipy.sparse.linalg.SuperLU 

373 Object, which has a ``solve`` method. 

374 

375 See also 

376 -------- 

377 spilu : incomplete LU decomposition 

378 

379 Notes 

380 ----- 

381 This function uses the SuperLU library. 

382 

383 References 

384 ---------- 

385 .. [1] SuperLU https://portal.nersc.gov/project/sparse/superlu/ 

386 

387 Examples 

388 -------- 

389 >>> import numpy as np 

390 >>> from scipy.sparse import csc_matrix 

391 >>> from scipy.sparse.linalg import splu 

392 >>> A = csc_matrix([[1., 0., 0.], [5., 0., 2.], [0., -1., 0.]], dtype=float) 

393 >>> B = splu(A) 

394 >>> x = np.array([1., 2., 3.], dtype=float) 

395 >>> B.solve(x) 

396 array([ 1. , -3. , -1.5]) 

397 >>> A.dot(B.solve(x)) 

398 array([ 1., 2., 3.]) 

399 >>> B.solve(A.dot(x)) 

400 array([ 1., 2., 3.]) 

401 """ 

402 

403 if is_pydata_spmatrix(A): 

404 def csc_construct_func(*a, cls=type(A)): 

405 return cls(csc_matrix(*a)) 

406 A = A.to_scipy_sparse().tocsc() 

407 else: 

408 csc_construct_func = csc_matrix 

409 

410 if not (issparse(A) and A.format == "csc"): 

411 A = csc_matrix(A) 

412 warn('splu converted its input to CSC format', SparseEfficiencyWarning) 

413 

414 # sum duplicates for non-canonical format 

415 A.sum_duplicates() 

416 A = A._asfptype() # upcast to a floating point format 

417 

418 M, N = A.shape 

419 if (M != N): 

420 raise ValueError("can only factor square matrices") # is this true? 

421 

422 indices, indptr = _safe_downcast_indices(A) 

423 

424 _options = dict(DiagPivotThresh=diag_pivot_thresh, ColPerm=permc_spec, 

425 PanelSize=panel_size, Relax=relax) 

426 if options is not None: 

427 _options.update(options) 

428 

429 # Ensure that no column permutations are applied 

430 if (_options["ColPerm"] == "NATURAL"): 

431 _options["SymmetricMode"] = True 

432 

433 return _superlu.gstrf(N, A.nnz, A.data, indices, indptr, 

434 csc_construct_func=csc_construct_func, 

435 ilu=False, options=_options) 

436 

437 

438def spilu(A, drop_tol=None, fill_factor=None, drop_rule=None, permc_spec=None, 

439 diag_pivot_thresh=None, relax=None, panel_size=None, options=None): 

440 """ 

441 Compute an incomplete LU decomposition for a sparse, square matrix. 

442 

443 The resulting object is an approximation to the inverse of `A`. 

444 

445 Parameters 

446 ---------- 

447 A : (N, N) array_like 

448 Sparse matrix to factorize. Most efficient when provided in CSC format. 

449 Other formats will be converted to CSC before factorization. 

450 drop_tol : float, optional 

451 Drop tolerance (0 <= tol <= 1) for an incomplete LU decomposition. 

452 (default: 1e-4) 

453 fill_factor : float, optional 

454 Specifies the fill ratio upper bound (>= 1.0) for ILU. (default: 10) 

455 drop_rule : str, optional 

456 Comma-separated string of drop rules to use. 

457 Available rules: ``basic``, ``prows``, ``column``, ``area``, 

458 ``secondary``, ``dynamic``, ``interp``. (Default: ``basic,area``) 

459 

460 See SuperLU documentation for details. 

461 

462 Remaining other options 

463 Same as for `splu` 

464 

465 Returns 

466 ------- 

467 invA_approx : scipy.sparse.linalg.SuperLU 

468 Object, which has a ``solve`` method. 

469 

470 See also 

471 -------- 

472 splu : complete LU decomposition 

473 

474 Notes 

475 ----- 

476 To improve the better approximation to the inverse, you may need to 

477 increase `fill_factor` AND decrease `drop_tol`. 

478 

479 This function uses the SuperLU library. 

480 

481 Examples 

482 -------- 

483 >>> import numpy as np 

484 >>> from scipy.sparse import csc_matrix 

485 >>> from scipy.sparse.linalg import spilu 

486 >>> A = csc_matrix([[1., 0., 0.], [5., 0., 2.], [0., -1., 0.]], dtype=float) 

487 >>> B = spilu(A) 

488 >>> x = np.array([1., 2., 3.], dtype=float) 

489 >>> B.solve(x) 

490 array([ 1. , -3. , -1.5]) 

491 >>> A.dot(B.solve(x)) 

492 array([ 1., 2., 3.]) 

493 >>> B.solve(A.dot(x)) 

494 array([ 1., 2., 3.]) 

495 """ 

496 

497 if is_pydata_spmatrix(A): 

498 def csc_construct_func(*a, cls=type(A)): 

499 return cls(csc_matrix(*a)) 

500 A = A.to_scipy_sparse().tocsc() 

501 else: 

502 csc_construct_func = csc_matrix 

503 

504 if not (issparse(A) and A.format == "csc"): 

505 A = csc_matrix(A) 

506 warn('spilu converted its input to CSC format', 

507 SparseEfficiencyWarning) 

508 

509 # sum duplicates for non-canonical format 

510 A.sum_duplicates() 

511 A = A._asfptype() # upcast to a floating point format 

512 

513 M, N = A.shape 

514 if (M != N): 

515 raise ValueError("can only factor square matrices") # is this true? 

516 

517 indices, indptr = _safe_downcast_indices(A) 

518 

519 _options = dict(ILU_DropRule=drop_rule, ILU_DropTol=drop_tol, 

520 ILU_FillFactor=fill_factor, 

521 DiagPivotThresh=diag_pivot_thresh, ColPerm=permc_spec, 

522 PanelSize=panel_size, Relax=relax) 

523 if options is not None: 

524 _options.update(options) 

525 

526 # Ensure that no column permutations are applied 

527 if (_options["ColPerm"] == "NATURAL"): 

528 _options["SymmetricMode"] = True 

529 

530 return _superlu.gstrf(N, A.nnz, A.data, indices, indptr, 

531 csc_construct_func=csc_construct_func, 

532 ilu=True, options=_options) 

533 

534 

535def factorized(A): 

536 """ 

537 Return a function for solving a sparse linear system, with A pre-factorized. 

538 

539 Parameters 

540 ---------- 

541 A : (N, N) array_like 

542 Input. A in CSC format is most efficient. A CSR format matrix will 

543 be converted to CSC before factorization. 

544 

545 Returns 

546 ------- 

547 solve : callable 

548 To solve the linear system of equations given in `A`, the `solve` 

549 callable should be passed an ndarray of shape (N,). 

550 

551 Examples 

552 -------- 

553 >>> import numpy as np 

554 >>> from scipy.sparse.linalg import factorized 

555 >>> A = np.array([[ 3. , 2. , -1. ], 

556 ... [ 2. , -2. , 4. ], 

557 ... [-1. , 0.5, -1. ]]) 

558 >>> solve = factorized(A) # Makes LU decomposition. 

559 >>> rhs1 = np.array([1, -2, 0]) 

560 >>> solve(rhs1) # Uses the LU factors. 

561 array([ 1., -2., -2.]) 

562 

563 """ 

564 if is_pydata_spmatrix(A): 

565 A = A.to_scipy_sparse().tocsc() 

566 

567 if useUmfpack: 

568 if noScikit: 

569 raise RuntimeError('Scikits.umfpack not installed.') 

570 

571 if not (issparse(A) and A.format == "csc"): 

572 A = csc_matrix(A) 

573 warn('splu converted its input to CSC format', 

574 SparseEfficiencyWarning) 

575 

576 A = A._asfptype() # upcast to a floating point format 

577 

578 if A.dtype.char not in 'dD': 

579 raise ValueError("convert matrix data to double, please, using" 

580 " .astype(), or set linsolve.useUmfpack = False") 

581 

582 umf_family, A = _get_umf_family(A) 

583 umf = umfpack.UmfpackContext(umf_family) 

584 

585 # Make LU decomposition. 

586 umf.numeric(A) 

587 

588 def solve(b): 

589 with np.errstate(divide="ignore", invalid="ignore"): 

590 # Ignoring warnings with numpy >= 1.23.0, see gh-16523 

591 result = umf.solve(umfpack.UMFPACK_A, A, b, autoTranspose=True) 

592 

593 return result 

594 

595 return solve 

596 else: 

597 return splu(A).solve 

598 

599 

600def spsolve_triangular(A, b, lower=True, overwrite_A=False, overwrite_b=False, 

601 unit_diagonal=False): 

602 """ 

603 Solve the equation ``A x = b`` for `x`, assuming A is a triangular matrix. 

604 

605 Parameters 

606 ---------- 

607 A : (M, M) sparse matrix 

608 A sparse square triangular matrix. Should be in CSR format. 

609 b : (M,) or (M, N) array_like 

610 Right-hand side matrix in ``A x = b`` 

611 lower : bool, optional 

612 Whether `A` is a lower or upper triangular matrix. 

613 Default is lower triangular matrix. 

614 overwrite_A : bool, optional 

615 Allow changing `A`. The indices of `A` are going to be sorted and zero 

616 entries are going to be removed. 

617 Enabling gives a performance gain. Default is False. 

618 overwrite_b : bool, optional 

619 Allow overwriting data in `b`. 

620 Enabling gives a performance gain. Default is False. 

621 If `overwrite_b` is True, it should be ensured that 

622 `b` has an appropriate dtype to be able to store the result. 

623 unit_diagonal : bool, optional 

624 If True, diagonal elements of `a` are assumed to be 1 and will not be 

625 referenced. 

626 

627 .. versionadded:: 1.4.0 

628 

629 Returns 

630 ------- 

631 x : (M,) or (M, N) ndarray 

632 Solution to the system ``A x = b``. Shape of return matches shape 

633 of `b`. 

634 

635 Raises 

636 ------ 

637 LinAlgError 

638 If `A` is singular or not triangular. 

639 ValueError 

640 If shape of `A` or shape of `b` do not match the requirements. 

641 

642 Notes 

643 ----- 

644 .. versionadded:: 0.19.0 

645 

646 Examples 

647 -------- 

648 >>> import numpy as np 

649 >>> from scipy.sparse import csr_matrix 

650 >>> from scipy.sparse.linalg import spsolve_triangular 

651 >>> A = csr_matrix([[3, 0, 0], [1, -1, 0], [2, 0, 1]], dtype=float) 

652 >>> B = np.array([[2, 0], [-1, 0], [2, 0]], dtype=float) 

653 >>> x = spsolve_triangular(A, B) 

654 >>> np.allclose(A.dot(x), B) 

655 True 

656 """ 

657 

658 if is_pydata_spmatrix(A): 

659 A = A.to_scipy_sparse().tocsr() 

660 

661 # Check the input for correct type and format. 

662 if not (issparse(A) and A.format == "csr"): 

663 warn('CSR matrix format is required. Converting to CSR matrix.', 

664 SparseEfficiencyWarning) 

665 A = csr_matrix(A) 

666 elif not overwrite_A: 

667 A = A.copy() 

668 

669 if A.shape[0] != A.shape[1]: 

670 raise ValueError( 

671 f'A must be a square matrix but its shape is {A.shape}.') 

672 

673 # sum duplicates for non-canonical format 

674 A.sum_duplicates() 

675 

676 b = np.asanyarray(b) 

677 

678 if b.ndim not in [1, 2]: 

679 raise ValueError( 

680 f'b must have 1 or 2 dims but its shape is {b.shape}.') 

681 if A.shape[0] != b.shape[0]: 

682 raise ValueError( 

683 'The size of the dimensions of A must be equal to ' 

684 'the size of the first dimension of b but the shape of A is ' 

685 '{} and the shape of b is {}.'.format(A.shape, b.shape)) 

686 

687 # Init x as (a copy of) b. 

688 x_dtype = np.result_type(A.data, b, np.float64) 

689 if overwrite_b: 

690 if np.can_cast(b.dtype, x_dtype, casting='same_kind'): 

691 x = b 

692 else: 

693 raise ValueError( 

694 'Cannot overwrite b (dtype {}) with result ' 

695 'of type {}.'.format(b.dtype, x_dtype)) 

696 else: 

697 x = b.astype(x_dtype, copy=True) 

698 

699 # Choose forward or backward order. 

700 if lower: 

701 row_indices = range(len(b)) 

702 else: 

703 row_indices = range(len(b) - 1, -1, -1) 

704 

705 # Fill x iteratively. 

706 for i in row_indices: 

707 

708 # Get indices for i-th row. 

709 indptr_start = A.indptr[i] 

710 indptr_stop = A.indptr[i + 1] 

711 

712 if lower: 

713 A_diagonal_index_row_i = indptr_stop - 1 

714 A_off_diagonal_indices_row_i = slice(indptr_start, indptr_stop - 1) 

715 else: 

716 A_diagonal_index_row_i = indptr_start 

717 A_off_diagonal_indices_row_i = slice(indptr_start + 1, indptr_stop) 

718 

719 # Check regularity and triangularity of A. 

720 if not unit_diagonal and (indptr_stop <= indptr_start 

721 or A.indices[A_diagonal_index_row_i] < i): 

722 raise LinAlgError( 

723 f'A is singular: diagonal {i} is zero.') 

724 if not unit_diagonal and A.indices[A_diagonal_index_row_i] > i: 

725 raise LinAlgError( 

726 'A is not triangular: A[{}, {}] is nonzero.' 

727 ''.format(i, A.indices[A_diagonal_index_row_i])) 

728 

729 # Incorporate off-diagonal entries. 

730 A_column_indices_in_row_i = A.indices[A_off_diagonal_indices_row_i] 

731 A_values_in_row_i = A.data[A_off_diagonal_indices_row_i] 

732 x[i] -= np.dot(x[A_column_indices_in_row_i].T, A_values_in_row_i) 

733 

734 # Compute i-th entry of x. 

735 if not unit_diagonal: 

736 x[i] /= A.data[A_diagonal_index_row_i] 

737 

738 return x