"""
Easy-to-use UMFPACK interface
=============================
.. currentmodule:: scikits.umfpack
The following functions can be used for LU decompositions and solving
equation systems:
.. autosummary::
:toctree: reference/
spsolve
splu
UmfpackLU
"""
from __future__ import division, print_function, absolute_import
from warnings import warn
import sys
import numpy as np
from numpy import asarray
from scipy.sparse import (isspmatrix_csc, isspmatrix_csr, isspmatrix,
SparseEfficiencyWarning, csc_matrix, hstack)
from .umfpack import UmfpackContext, UMFPACK_A
_families = {
(np.float64, np.int32): 'di',
(np.complex128, np.int32): 'zi',
(np.float64, np.int64): 'dl',
(np.complex128, np.int64): 'zl'
}
__all__ = ['spsolve', 'splu', 'UmfpackLU']
if sys.version_info[0] >= 3:
xrange = range
def spsolve(A, b):
"""Solve the sparse linear system Ax=b, where b may be a vector or a matrix.
Parameters
----------
A : ndarray or sparse matrix
The square matrix A will be converted into CSC or CSR form
b : ndarray or sparse matrix
The matrix or vector representing the right hand side of the equation.
Returns
-------
x : ndarray or sparse matrix
the solution of the sparse linear equation.
If b is a vector, then x is a vector of size A.shape[0]
If b is a matrix, then x is a matrix of size (A.shape[0],)+b.shape[1:]
"""
x = UmfpackLU(A).solve(b)
if b.ndim == 2 and b.shape[1] == 1:
# compatibility with scipy.sparse.spsolve quirk
return x.ravel()
else:
return x
def splu(A):
"""
Compute the LU decomposition of a sparse, square matrix.
Parameters
----------
A : sparse matrix
Sparse matrix to factorize. Should be in CSR or CSC format.
Returns
-------
invA : scikits.umfpack.UmfpackLU
Object, which has a ``solve`` method.
Notes
-----
This function uses the UMFPACK library.
"""
return UmfpackLU(A)
class UmfpackLU(object):
"""
LU factorization of a sparse matrix.
Factorization is represented as::
Pr * (R^-1) * A * Pc = L * U
Parameters
----------
A : csc_matrix or csr_matrix
Matrix to decompose
Attributes
----------
shape
nnz
perm_c
perm_r
L
U
R
Methods
-------
solve
solve_sparse
Examples
--------
The LU decomposition can be used to solve matrix equations. Consider:
>>> import numpy as np
>>> from scipy.sparse import csc_matrix
>>> from scikits import umfpack
>>> A = csc_matrix([[1,2,0,4],[1,0,0,1],[1,0,2,1],[2,2,1,0.]])
This can be solved for a given right-hand side:
>>> lu = umfpack.splu(A)
>>> b = np.array([1, 2, 3, 4])
>>> x = lu.solve(b)
>>> A.dot(x)
array([ 1., 2., 3., 4.])
The ``lu`` object also contains an explicit representation of the
decomposition. The permutations are represented as mappings of
indices:
>>> lu.perm_r
array([0, 2, 1, 3], dtype=int32)
>>> lu.perm_c
array([2, 0, 1, 3], dtype=int32)
The L and U factors are sparse matrices in CSC format:
>>> lu.L.A
array([[ 1. , 0. , 0. , 0. ],
[ 0. , 1. , 0. , 0. ],
[ 0. , 0. , 1. , 0. ],
[ 1. , 0.5, 0.5, 1. ]])
>>> lu.U.A
array([[ 2., 0., 1., 4.],
[ 0., 2., 1., 1.],
[ 0., 0., 1., 1.],
[ 0., 0., 0., -5.]])
The permutation matrices can be constructed:
>>> Pr = csc_matrix((4, 4))
>>> Pr[lu.perm_r, np.arange(4)] = 1
>>> Pc = csc_matrix((4, 4))
>>> Pc[np.arange(4), lu.perm_c] = 1
Similarly for the row scalings:
>>> R = csc_matrix((4, 4))
>>> R.setdiag(lu.R)
We can reassemble the original matrix:
>>> (Pr.T * R * (lu.L * lu.U) * Pc.T).A
array([[ 1., 2., 0., 4.],
[ 1., 0., 0., 1.],
[ 1., 0., 2., 1.],
[ 2., 2., 1., 0.]])
"""
def __init__(self, A):
if not (isspmatrix_csc(A) or isspmatrix_csr(A)):
A = csc_matrix(A)
warn('spsolve requires A be CSC or CSR matrix format',
SparseEfficiencyWarning)
A.sort_indices()
A = A.asfptype() # upcast to a floating point format
M, N = A.shape
if (M != N):
raise ValueError("matrix must be square (has shape %s)" % ((M, N),))
f_type = np.sctypeDict[A.dtype.name]
i_type = np.sctypeDict[A.indices.dtype.name]
try:
family = _families[(f_type, i_type)]
except KeyError:
msg = 'only float64 or complex128 matrices with int32 or int64' \
' indices are supported! (got: matrix: %s, indices: %s)' \
% (f_type, i_type)
raise ValueError(msg)
self.umf = UmfpackContext(family)
self.umf.numeric(A)
self._A = A
self._L = None
self._U = None
self._P = None
self._Q = None
self._R = None
def solve(self, b):
"""
Solve linear equation A x = b for x
Parameters
----------
b : ndarray
Right-hand side of the matrix equation. Can be vector or a matrix.
Returns
-------
x : ndarray
Solution to the matrix equation
"""
if isspmatrix(b):
b = b.toarray()
if b.shape[0] != self._A.shape[1]:
raise ValueError("Shape of b is not compatible with that of A")
b_arr = asarray(b, dtype=self._A.dtype).reshape(b.shape[0], -1)
x = np.zeros((self._A.shape[0], b_arr.shape[1]), dtype=self._A.dtype)
for j in range(b_arr.shape[1]):
x[:,j] = self.umf.solve(UMFPACK_A, self._A, b_arr[:,j], autoTranspose=True)
return x.reshape((self._A.shape[0],) + b.shape[1:])
def solve_sparse(self, B):
"""
Solve linear equation of the form A X = B. Where B and X are sparse matrices.
Parameters
----------
B : any scipy.sparse matrix
Right-hand side of the matrix equation.
Note: it will be converted to csc_matrix via `.tocsc()`.
Returns
-------
X : csc_matrix
Solution to the matrix equation as a csc_matrix
"""
B = B.tocsc()
cols = list()
for j in xrange(B.shape[1]):
col = self.solve(B[:,j])
cols.append(csc_matrix(col))
return hstack(cols)
def _compute_lu(self):
if self._L is None:
self._L, self._U, self._P, self._Q, self._R, do_recip = self.umf.lu(self._A)
if do_recip:
with np.errstate(divide='ignore'):
np.reciprocal(self._R, out=self._R)
# Conform to scipy.sparse.splu convention on permutation matrices
self._P = self._P[self._P]
@property
def shape(self):
"""
Shape of the original matrix as a tuple of ints.
"""
return self._A.shape
@property
def L(self):
"""
Lower triangular factor with unit diagonal as a
`scipy.sparse.csc_matrix`.
"""
self._compute_lu()
return self._L
@property
def U(self):
"""
Upper triangular factor as a `scipy.sparse.csc_matrix`.
"""
self._compute_lu()
return self._U
@property
def R(self):
"""
Row scaling factors, as a 1D array.
"""
self._compute_lu()
return self._R
@property
def perm_c(self):
"""
Permutation Pc represented as an array of indices.
The column permutation matrix can be reconstructed via:
>>> Pc = np.zeros((n, n))
>>> Pc[np.arange(n), perm_c] = 1
"""
self._compute_lu()
return self._Q
@property
def perm_r(self):
"""
Permutation Pr represented as an array of indices.
The row permutation matrix can be reconstructed via:
>>> Pr = np.zeros((n, n))
>>> Pr[perm_r, np.arange(n)] = 1
"""
self._compute_lu()
return self._P
@property
def nnz(self):
"""
Combined number of nonzeros in L and U: L.nnz + U.nnz
"""
return self._L.nnz + self._U.nnz
