Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.9/dist-packages/scipy/sparse/linalg/_dsolve/linsolve.py: 11%
225 statements
« prev ^ index » next coverage.py v7.3.1, created at 2023-09-23 06:43 +0000
« prev ^ index » next coverage.py v7.3.1, created at 2023-09-23 06:43 +0000
1from warnings import warn
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
11from . import _superlu
13noScikit = False
14try:
15 import scikits.umfpack as umfpack
16except ImportError:
17 noScikit = True
19useUmfpack = not noScikit
21__all__ = ['use_solver', 'spsolve', 'splu', 'spilu', 'factorized',
22 'MatrixRankWarning', 'spsolve_triangular']
25class MatrixRankWarning(UserWarning):
26 pass
29def use_solver(**kwargs):
30 """
31 Select default sparse direct solver to be used.
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
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.
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.
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
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
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
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.
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'])
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 }
104 # A.dtype.name can only be "float64" or
105 # "complex128" in control flow
106 f_type = getattr(np, A.dtype.name)
107 # control flow may allow for more index
108 # types to get through here
109 i_type = getattr(np, A.indices.dtype.name)
111 try:
112 family = _families[(f_type, i_type)]
114 except KeyError as e:
115 msg = 'only float64 or complex128 matrices with int32 or int64' \
116 ' indices are supported! (got: matrix: %s, indices: %s)' \
117 % (f_type, i_type)
118 raise ValueError(msg) from e
120 # See gh-8278. Considered converting only if
121 # A.shape[0]*A.shape[1] > np.iinfo(np.int32).max,
122 # but that didn't always fix the issue.
123 family = family[0] + "l"
124 A_new = copy.copy(A)
125 A_new.indptr = np.array(A.indptr, copy=False, dtype=np.int64)
126 A_new.indices = np.array(A.indices, copy=False, dtype=np.int64)
128 return family, A_new
130def _safe_downcast_indices(A):
131 # check for safe downcasting
132 max_value = np.iinfo(np.intc).max
134 if A.indptr[-1] > max_value: # indptr[-1] is max b/c indptr always sorted
135 raise ValueError("indptr values too large for SuperLU")
137 if max(*A.shape) > max_value: # only check large enough arrays
138 if np.any(A.indices > max_value):
139 raise ValueError("indices values too large for SuperLU")
141 indices = A.indices.astype(np.intc, copy=False)
142 indptr = A.indptr.astype(np.intc, copy=False)
143 return indices, indptr
145def spsolve(A, b, permc_spec=None, use_umfpack=True):
146 """Solve the sparse linear system Ax=b, where b may be a vector or a matrix.
148 Parameters
149 ----------
150 A : ndarray or sparse matrix
151 The square matrix A will be converted into CSC or CSR form
152 b : ndarray or sparse matrix
153 The matrix or vector representing the right hand side of the equation.
154 If a vector, b.shape must be (n,) or (n, 1).
155 permc_spec : str, optional
156 How to permute the columns of the matrix for sparsity preservation.
157 (default: 'COLAMD')
159 - ``NATURAL``: natural ordering.
160 - ``MMD_ATA``: minimum degree ordering on the structure of A^T A.
161 - ``MMD_AT_PLUS_A``: minimum degree ordering on the structure of A^T+A.
162 - ``COLAMD``: approximate minimum degree column ordering [1]_, [2]_.
164 use_umfpack : bool, optional
165 if True (default) then use UMFPACK for the solution [3]_, [4]_, [5]_,
166 [6]_ . This is only referenced if b is a vector and
167 ``scikits.umfpack`` is installed.
169 Returns
170 -------
171 x : ndarray or sparse matrix
172 the solution of the sparse linear equation.
173 If b is a vector, then x is a vector of size A.shape[1]
174 If b is a matrix, then x is a matrix of size (A.shape[1], b.shape[1])
176 Notes
177 -----
178 For solving the matrix expression AX = B, this solver assumes the resulting
179 matrix X is sparse, as is often the case for very sparse inputs. If the
180 resulting X is dense, the construction of this sparse result will be
181 relatively expensive. In that case, consider converting A to a dense
182 matrix and using scipy.linalg.solve or its variants.
184 References
185 ----------
186 .. [1] T. A. Davis, J. R. Gilbert, S. Larimore, E. Ng, Algorithm 836:
187 COLAMD, an approximate column minimum degree ordering algorithm,
188 ACM Trans. on Mathematical Software, 30(3), 2004, pp. 377--380.
189 :doi:`10.1145/1024074.1024080`
191 .. [2] T. A. Davis, J. R. Gilbert, S. Larimore, E. Ng, A column approximate
192 minimum degree ordering algorithm, ACM Trans. on Mathematical
193 Software, 30(3), 2004, pp. 353--376. :doi:`10.1145/1024074.1024079`
195 .. [3] T. A. Davis, Algorithm 832: UMFPACK - an unsymmetric-pattern
196 multifrontal method with a column pre-ordering strategy, ACM
197 Trans. on Mathematical Software, 30(2), 2004, pp. 196--199.
198 https://dl.acm.org/doi/abs/10.1145/992200.992206
200 .. [4] T. A. Davis, A column pre-ordering strategy for the
201 unsymmetric-pattern multifrontal method, ACM Trans.
202 on Mathematical Software, 30(2), 2004, pp. 165--195.
203 https://dl.acm.org/doi/abs/10.1145/992200.992205
205 .. [5] T. A. Davis and I. S. Duff, A combined unifrontal/multifrontal
206 method for unsymmetric sparse matrices, ACM Trans. on
207 Mathematical Software, 25(1), 1999, pp. 1--19.
208 https://doi.org/10.1145/305658.287640
210 .. [6] T. A. Davis and I. S. Duff, An unsymmetric-pattern multifrontal
211 method for sparse LU factorization, SIAM J. Matrix Analysis and
212 Computations, 18(1), 1997, pp. 140--158.
213 https://doi.org/10.1137/S0895479894246905T.
216 Examples
217 --------
218 >>> import numpy as np
219 >>> from scipy.sparse import csc_matrix
220 >>> from scipy.sparse.linalg import spsolve
221 >>> A = csc_matrix([[3, 2, 0], [1, -1, 0], [0, 5, 1]], dtype=float)
222 >>> B = csc_matrix([[2, 0], [-1, 0], [2, 0]], dtype=float)
223 >>> x = spsolve(A, B)
224 >>> np.allclose(A.dot(x).toarray(), B.toarray())
225 True
226 """
228 if is_pydata_spmatrix(A):
229 A = A.to_scipy_sparse().tocsc()
231 if not (issparse(A) and A.format in ("csc", "csr")):
232 A = csc_matrix(A)
233 warn('spsolve requires A be CSC or CSR matrix format',
234 SparseEfficiencyWarning)
236 # b is a vector only if b have shape (n,) or (n, 1)
237 b_is_sparse = issparse(b) or is_pydata_spmatrix(b)
238 if not b_is_sparse:
239 b = asarray(b)
240 b_is_vector = ((b.ndim == 1) or (b.ndim == 2 and b.shape[1] == 1))
242 # sum duplicates for non-canonical format
243 A.sum_duplicates()
244 A = A._asfptype() # upcast to a floating point format
245 result_dtype = np.promote_types(A.dtype, b.dtype)
246 if A.dtype != result_dtype:
247 A = A.astype(result_dtype)
248 if b.dtype != result_dtype:
249 b = b.astype(result_dtype)
251 # validate input shapes
252 M, N = A.shape
253 if (M != N):
254 raise ValueError(f"matrix must be square (has shape {(M, N)})")
256 if M != b.shape[0]:
257 raise ValueError("matrix - rhs dimension mismatch (%s - %s)"
258 % (A.shape, b.shape[0]))
260 use_umfpack = use_umfpack and useUmfpack
262 if b_is_vector and use_umfpack:
263 if b_is_sparse:
264 b_vec = b.toarray()
265 else:
266 b_vec = b
267 b_vec = asarray(b_vec, dtype=A.dtype).ravel()
269 if noScikit:
270 raise RuntimeError('Scikits.umfpack not installed.')
272 if A.dtype.char not in 'dD':
273 raise ValueError("convert matrix data to double, please, using"
274 " .astype(), or set linsolve.useUmfpack = False")
276 umf_family, A = _get_umf_family(A)
277 umf = umfpack.UmfpackContext(umf_family)
278 x = umf.linsolve(umfpack.UMFPACK_A, A, b_vec,
279 autoTranspose=True)
280 else:
281 if b_is_vector and b_is_sparse:
282 b = b.toarray()
283 b_is_sparse = False
285 if not b_is_sparse:
286 if A.format == "csc":
287 flag = 1 # CSC format
288 else:
289 flag = 0 # CSR format
291 indices = A.indices.astype(np.intc, copy=False)
292 indptr = A.indptr.astype(np.intc, copy=False)
293 options = dict(ColPerm=permc_spec)
294 x, info = _superlu.gssv(N, A.nnz, A.data, indices, indptr,
295 b, flag, options=options)
296 if info != 0:
297 warn("Matrix is exactly singular", MatrixRankWarning)
298 x.fill(np.nan)
299 if b_is_vector:
300 x = x.ravel()
301 else:
302 # b is sparse
303 Afactsolve = factorized(A)
305 if not (b.format == "csc" or is_pydata_spmatrix(b)):
306 warn('spsolve is more efficient when sparse b '
307 'is in the CSC matrix format', SparseEfficiencyWarning)
308 b = csc_matrix(b)
310 # Create a sparse output matrix by repeatedly applying
311 # the sparse factorization to solve columns of b.
312 data_segs = []
313 row_segs = []
314 col_segs = []
315 for j in range(b.shape[1]):
316 # TODO: replace this with
317 # bj = b[:, j].toarray().ravel()
318 # once 1D sparse arrays are supported.
319 # That is a slightly faster code path.
320 bj = b[:, [j]].toarray().ravel()
321 xj = Afactsolve(bj)
322 w = np.flatnonzero(xj)
323 segment_length = w.shape[0]
324 row_segs.append(w)
325 col_segs.append(np.full(segment_length, j, dtype=int))
326 data_segs.append(np.asarray(xj[w], dtype=A.dtype))
327 sparse_data = np.concatenate(data_segs)
328 sparse_row = np.concatenate(row_segs)
329 sparse_col = np.concatenate(col_segs)
330 x = A.__class__((sparse_data, (sparse_row, sparse_col)),
331 shape=b.shape, dtype=A.dtype)
333 if is_pydata_spmatrix(b):
334 x = b.__class__(x)
336 return x
339def splu(A, permc_spec=None, diag_pivot_thresh=None,
340 relax=None, panel_size=None, options=dict()):
341 """
342 Compute the LU decomposition of a sparse, square matrix.
344 Parameters
345 ----------
346 A : sparse matrix
347 Sparse matrix to factorize. Most efficient when provided in CSC
348 format. Other formats will be converted to CSC before factorization.
349 permc_spec : str, optional
350 How to permute the columns of the matrix for sparsity preservation.
351 (default: 'COLAMD')
353 - ``NATURAL``: natural ordering.
354 - ``MMD_ATA``: minimum degree ordering on the structure of A^T A.
355 - ``MMD_AT_PLUS_A``: minimum degree ordering on the structure of A^T+A.
356 - ``COLAMD``: approximate minimum degree column ordering
358 diag_pivot_thresh : float, optional
359 Threshold used for a diagonal entry to be an acceptable pivot.
360 See SuperLU user's guide for details [1]_
361 relax : int, optional
362 Expert option for customizing the degree of relaxing supernodes.
363 See SuperLU user's guide for details [1]_
364 panel_size : int, optional
365 Expert option for customizing the panel size.
366 See SuperLU user's guide for details [1]_
367 options : dict, optional
368 Dictionary containing additional expert options to SuperLU.
369 See SuperLU user guide [1]_ (section 2.4 on the 'Options' argument)
370 for more details. For example, you can specify
371 ``options=dict(Equil=False, IterRefine='SINGLE'))``
372 to turn equilibration off and perform a single iterative refinement.
374 Returns
375 -------
376 invA : scipy.sparse.linalg.SuperLU
377 Object, which has a ``solve`` method.
379 See also
380 --------
381 spilu : incomplete LU decomposition
383 Notes
384 -----
385 This function uses the SuperLU library.
387 References
388 ----------
389 .. [1] SuperLU https://portal.nersc.gov/project/sparse/superlu/
391 Examples
392 --------
393 >>> import numpy as np
394 >>> from scipy.sparse import csc_matrix
395 >>> from scipy.sparse.linalg import splu
396 >>> A = csc_matrix([[1., 0., 0.], [5., 0., 2.], [0., -1., 0.]], dtype=float)
397 >>> B = splu(A)
398 >>> x = np.array([1., 2., 3.], dtype=float)
399 >>> B.solve(x)
400 array([ 1. , -3. , -1.5])
401 >>> A.dot(B.solve(x))
402 array([ 1., 2., 3.])
403 >>> B.solve(A.dot(x))
404 array([ 1., 2., 3.])
405 """
407 if is_pydata_spmatrix(A):
408 def csc_construct_func(*a, cls=type(A)):
409 return cls(csc_matrix(*a))
410 A = A.to_scipy_sparse().tocsc()
411 else:
412 csc_construct_func = csc_matrix
414 if not (issparse(A) and A.format == "csc"):
415 A = csc_matrix(A)
416 warn('splu converted its input to CSC format', SparseEfficiencyWarning)
418 # sum duplicates for non-canonical format
419 A.sum_duplicates()
420 A = A._asfptype() # upcast to a floating point format
422 M, N = A.shape
423 if (M != N):
424 raise ValueError("can only factor square matrices") # is this true?
426 indices, indptr = _safe_downcast_indices(A)
428 _options = dict(DiagPivotThresh=diag_pivot_thresh, ColPerm=permc_spec,
429 PanelSize=panel_size, Relax=relax)
430 if options is not None:
431 _options.update(options)
433 # Ensure that no column permutations are applied
434 if (_options["ColPerm"] == "NATURAL"):
435 _options["SymmetricMode"] = True
437 return _superlu.gstrf(N, A.nnz, A.data, indices, indptr,
438 csc_construct_func=csc_construct_func,
439 ilu=False, options=_options)
442def spilu(A, drop_tol=None, fill_factor=None, drop_rule=None, permc_spec=None,
443 diag_pivot_thresh=None, relax=None, panel_size=None, options=None):
444 """
445 Compute an incomplete LU decomposition for a sparse, square matrix.
447 The resulting object is an approximation to the inverse of `A`.
449 Parameters
450 ----------
451 A : (N, N) array_like
452 Sparse matrix to factorize. Most efficient when provided in CSC format.
453 Other formats will be converted to CSC before factorization.
454 drop_tol : float, optional
455 Drop tolerance (0 <= tol <= 1) for an incomplete LU decomposition.
456 (default: 1e-4)
457 fill_factor : float, optional
458 Specifies the fill ratio upper bound (>= 1.0) for ILU. (default: 10)
459 drop_rule : str, optional
460 Comma-separated string of drop rules to use.
461 Available rules: ``basic``, ``prows``, ``column``, ``area``,
462 ``secondary``, ``dynamic``, ``interp``. (Default: ``basic,area``)
464 See SuperLU documentation for details.
466 Remaining other options
467 Same as for `splu`
469 Returns
470 -------
471 invA_approx : scipy.sparse.linalg.SuperLU
472 Object, which has a ``solve`` method.
474 See also
475 --------
476 splu : complete LU decomposition
478 Notes
479 -----
480 To improve the better approximation to the inverse, you may need to
481 increase `fill_factor` AND decrease `drop_tol`.
483 This function uses the SuperLU library.
485 Examples
486 --------
487 >>> import numpy as np
488 >>> from scipy.sparse import csc_matrix
489 >>> from scipy.sparse.linalg import spilu
490 >>> A = csc_matrix([[1., 0., 0.], [5., 0., 2.], [0., -1., 0.]], dtype=float)
491 >>> B = spilu(A)
492 >>> x = np.array([1., 2., 3.], dtype=float)
493 >>> B.solve(x)
494 array([ 1. , -3. , -1.5])
495 >>> A.dot(B.solve(x))
496 array([ 1., 2., 3.])
497 >>> B.solve(A.dot(x))
498 array([ 1., 2., 3.])
499 """
501 if is_pydata_spmatrix(A):
502 def csc_construct_func(*a, cls=type(A)):
503 return cls(csc_matrix(*a))
504 A = A.to_scipy_sparse().tocsc()
505 else:
506 csc_construct_func = csc_matrix
508 if not (issparse(A) and A.format == "csc"):
509 A = csc_matrix(A)
510 warn('spilu converted its input to CSC format',
511 SparseEfficiencyWarning)
513 # sum duplicates for non-canonical format
514 A.sum_duplicates()
515 A = A._asfptype() # upcast to a floating point format
517 M, N = A.shape
518 if (M != N):
519 raise ValueError("can only factor square matrices") # is this true?
521 indices, indptr = _safe_downcast_indices(A)
523 _options = dict(ILU_DropRule=drop_rule, ILU_DropTol=drop_tol,
524 ILU_FillFactor=fill_factor,
525 DiagPivotThresh=diag_pivot_thresh, ColPerm=permc_spec,
526 PanelSize=panel_size, Relax=relax)
527 if options is not None:
528 _options.update(options)
530 # Ensure that no column permutations are applied
531 if (_options["ColPerm"] == "NATURAL"):
532 _options["SymmetricMode"] = True
534 return _superlu.gstrf(N, A.nnz, A.data, indices, indptr,
535 csc_construct_func=csc_construct_func,
536 ilu=True, options=_options)
539def factorized(A):
540 """
541 Return a function for solving a sparse linear system, with A pre-factorized.
543 Parameters
544 ----------
545 A : (N, N) array_like
546 Input. A in CSC format is most efficient. A CSR format matrix will
547 be converted to CSC before factorization.
549 Returns
550 -------
551 solve : callable
552 To solve the linear system of equations given in `A`, the `solve`
553 callable should be passed an ndarray of shape (N,).
555 Examples
556 --------
557 >>> import numpy as np
558 >>> from scipy.sparse.linalg import factorized
559 >>> from scipy.sparse import csc_matrix
560 >>> A = np.array([[ 3. , 2. , -1. ],
561 ... [ 2. , -2. , 4. ],
562 ... [-1. , 0.5, -1. ]])
563 >>> solve = factorized(csc_matrix(A)) # Makes LU decomposition.
564 >>> rhs1 = np.array([1, -2, 0])
565 >>> solve(rhs1) # Uses the LU factors.
566 array([ 1., -2., -2.])
568 """
569 if is_pydata_spmatrix(A):
570 A = A.to_scipy_sparse().tocsc()
572 if useUmfpack:
573 if noScikit:
574 raise RuntimeError('Scikits.umfpack not installed.')
576 if not (issparse(A) and A.format == "csc"):
577 A = csc_matrix(A)
578 warn('splu converted its input to CSC format',
579 SparseEfficiencyWarning)
581 A = A._asfptype() # upcast to a floating point format
583 if A.dtype.char not in 'dD':
584 raise ValueError("convert matrix data to double, please, using"
585 " .astype(), or set linsolve.useUmfpack = False")
587 umf_family, A = _get_umf_family(A)
588 umf = umfpack.UmfpackContext(umf_family)
590 # Make LU decomposition.
591 umf.numeric(A)
593 def solve(b):
594 with np.errstate(divide="ignore", invalid="ignore"):
595 # Ignoring warnings with numpy >= 1.23.0, see gh-16523
596 result = umf.solve(umfpack.UMFPACK_A, A, b, autoTranspose=True)
598 return result
600 return solve
601 else:
602 return splu(A).solve
605def spsolve_triangular(A, b, lower=True, overwrite_A=False, overwrite_b=False,
606 unit_diagonal=False):
607 """
608 Solve the equation ``A x = b`` for `x`, assuming A is a triangular matrix.
610 Parameters
611 ----------
612 A : (M, M) sparse matrix
613 A sparse square triangular matrix. Should be in CSR format.
614 b : (M,) or (M, N) array_like
615 Right-hand side matrix in ``A x = b``
616 lower : bool, optional
617 Whether `A` is a lower or upper triangular matrix.
618 Default is lower triangular matrix.
619 overwrite_A : bool, optional
620 Allow changing `A`. The indices of `A` are going to be sorted and zero
621 entries are going to be removed.
622 Enabling gives a performance gain. Default is False.
623 overwrite_b : bool, optional
624 Allow overwriting data in `b`.
625 Enabling gives a performance gain. Default is False.
626 If `overwrite_b` is True, it should be ensured that
627 `b` has an appropriate dtype to be able to store the result.
628 unit_diagonal : bool, optional
629 If True, diagonal elements of `a` are assumed to be 1 and will not be
630 referenced.
632 .. versionadded:: 1.4.0
634 Returns
635 -------
636 x : (M,) or (M, N) ndarray
637 Solution to the system ``A x = b``. Shape of return matches shape
638 of `b`.
640 Raises
641 ------
642 LinAlgError
643 If `A` is singular or not triangular.
644 ValueError
645 If shape of `A` or shape of `b` do not match the requirements.
647 Notes
648 -----
649 .. versionadded:: 0.19.0
651 Examples
652 --------
653 >>> import numpy as np
654 >>> from scipy.sparse import csr_matrix
655 >>> from scipy.sparse.linalg import spsolve_triangular
656 >>> A = csr_matrix([[3, 0, 0], [1, -1, 0], [2, 0, 1]], dtype=float)
657 >>> B = np.array([[2, 0], [-1, 0], [2, 0]], dtype=float)
658 >>> x = spsolve_triangular(A, B)
659 >>> np.allclose(A.dot(x), B)
660 True
661 """
663 if is_pydata_spmatrix(A):
664 A = A.to_scipy_sparse().tocsr()
666 # Check the input for correct type and format.
667 if not (issparse(A) and A.format == "csr"):
668 warn('CSR matrix format is required. Converting to CSR matrix.',
669 SparseEfficiencyWarning)
670 A = csr_matrix(A)
671 elif not overwrite_A:
672 A = A.copy()
674 if A.shape[0] != A.shape[1]:
675 raise ValueError(
676 f'A must be a square matrix but its shape is {A.shape}.')
678 # sum duplicates for non-canonical format
679 A.sum_duplicates()
681 b = np.asanyarray(b)
683 if b.ndim not in [1, 2]:
684 raise ValueError(
685 f'b must have 1 or 2 dims but its shape is {b.shape}.')
686 if A.shape[0] != b.shape[0]:
687 raise ValueError(
688 'The size of the dimensions of A must be equal to '
689 'the size of the first dimension of b but the shape of A is '
690 '{} and the shape of b is {}.'.format(A.shape, b.shape))
692 # Init x as (a copy of) b.
693 x_dtype = np.result_type(A.data, b, np.float64)
694 if overwrite_b:
695 if np.can_cast(b.dtype, x_dtype, casting='same_kind'):
696 x = b
697 else:
698 raise ValueError(
699 'Cannot overwrite b (dtype {}) with result '
700 'of type {}.'.format(b.dtype, x_dtype))
701 else:
702 x = b.astype(x_dtype, copy=True)
704 # Choose forward or backward order.
705 if lower:
706 row_indices = range(len(b))
707 else:
708 row_indices = range(len(b) - 1, -1, -1)
710 # Fill x iteratively.
711 for i in row_indices:
713 # Get indices for i-th row.
714 indptr_start = A.indptr[i]
715 indptr_stop = A.indptr[i + 1]
717 if lower:
718 A_diagonal_index_row_i = indptr_stop - 1
719 A_off_diagonal_indices_row_i = slice(indptr_start, indptr_stop - 1)
720 else:
721 A_diagonal_index_row_i = indptr_start
722 A_off_diagonal_indices_row_i = slice(indptr_start + 1, indptr_stop)
724 # Check regularity and triangularity of A.
725 if not unit_diagonal and (indptr_stop <= indptr_start
726 or A.indices[A_diagonal_index_row_i] < i):
727 raise LinAlgError(
728 f'A is singular: diagonal {i} is zero.')
729 if not unit_diagonal and A.indices[A_diagonal_index_row_i] > i:
730 raise LinAlgError(
731 'A is not triangular: A[{}, {}] is nonzero.'
732 ''.format(i, A.indices[A_diagonal_index_row_i]))
734 # Incorporate off-diagonal entries.
735 A_column_indices_in_row_i = A.indices[A_off_diagonal_indices_row_i]
736 A_values_in_row_i = A.data[A_off_diagonal_indices_row_i]
737 x[i] -= np.dot(x[A_column_indices_in_row_i].T, A_values_in_row_i)
739 # Compute i-th entry of x.
740 if not unit_diagonal:
741 x[i] /= A.data[A_diagonal_index_row_i]
743 return x