mirror of
https://github.com/wassname/simpeg.git
synced 2026-06-29 02:47:34 +08:00
Rowan Cockett
260 lines
8.5 KiB
Python
260 lines
8.5 KiB
Python
import numpy as np
|
|
import scipy.sparse as sparse
|
|
import scipy.sparse.linalg as linalg
|
|
from SimPEG.utils import mkvc
|
|
|
|
DEFAULTS = {'direct':'scipy', 'forward':'fortran', 'backward':'fortran', 'diagonal':'python'}
|
|
|
|
try:
|
|
import TriSolve
|
|
except Exception, e:
|
|
try:
|
|
import os
|
|
# Note: this may not work from SublimeText, if that is the case, just run the command in your shell.
|
|
os.system('f2py -c TriSolve.f -m TriSolve')
|
|
import TriSolve
|
|
except Exception, e:
|
|
print 'Warning: Python backend is being used for solver. Run setup.py from the command line.'
|
|
DEFAULTS['forward'] = 'python'
|
|
DEFAULTS['backward'] = 'python'
|
|
|
|
class Solver(object):
|
|
"""
|
|
Solver is a light wrapper on the various types of
|
|
linear solvers available in python.
|
|
|
|
:param scipy.sparse A: Matrix
|
|
:param bool doDirect: if you want a direct solver
|
|
:param string flag: Matrix type flag for special solves: [None, 'L', 'U', 'D']
|
|
:param dict options: options which are passed to each sub solver, see each for details.
|
|
:rtype: Solver
|
|
:return: Solver
|
|
|
|
To use for direct solvers::
|
|
|
|
solve = Solver(A, doDirect=True, flag=None, options={'factorize':True,'backend':'scipy'})
|
|
x = solve.solve(rhs)
|
|
|
|
Or in one line::
|
|
|
|
x = Solver(A).solve(rhs)
|
|
|
|
The flag can be set to None, 'L', 'U', or 'D', for general, lower, upper, and diagonal matrices, respectively.
|
|
|
|
"""
|
|
def __init__(self, A, doDirect=True, flag=None, options={}):
|
|
assert type(doDirect) is bool, 'doDirect must be a boolean'
|
|
assert flag in [None, 'L', 'U', 'D'], "flag must be set to None, 'L', 'U', or 'D'"
|
|
|
|
self.A = A
|
|
|
|
self.dsolve = None
|
|
self.doDirect = doDirect
|
|
self.flag = flag
|
|
self.options = options
|
|
|
|
def solve(self, b):
|
|
"""
|
|
Solves the linear system.
|
|
|
|
.. math::
|
|
|
|
Ax=b
|
|
|
|
:param numpy.ndarray b: the right hand side
|
|
:rtype: numpy.ndarray
|
|
:return: x
|
|
"""
|
|
if self.flag is None and self.doDirect:
|
|
return self.solveDirect(b, **self.options)
|
|
elif self.flag is None and not self.doDirect:
|
|
return self.solveIter(b, **self.options)
|
|
elif self.flag == 'U':
|
|
return self.solveBackward(b, **self.options)
|
|
elif self.flag == 'L':
|
|
return self.solveForward(b, **self.options)
|
|
elif self.flag == 'D':
|
|
return self.solveDiagonal(b, **self.options)
|
|
else:
|
|
raise Exception('Unknown flag.')
|
|
pass
|
|
|
|
def clean(self):
|
|
"""Cleans up the memory"""
|
|
del self.dsolve
|
|
self.dsolve = None
|
|
|
|
def solveDirect(self, b, factorize=False, backend=None):
|
|
"""
|
|
Use solve instead of this interface.
|
|
|
|
:param bool factorize: if you want to factorize and store factors
|
|
:param str backend: which backend to use. Default is scipy
|
|
:rtype: numpy.ndarray
|
|
:return: x
|
|
"""
|
|
if backend is None: backend = DEFAULTS['direct']
|
|
|
|
assert np.shape(self.A)[1] == np.shape(b)[0], 'Dimension mismatch'
|
|
|
|
if factorize and self.dsolve is None:
|
|
self.A = self.A.tocsc() # for efficiency
|
|
self.dsolve = linalg.factorized(self.A)
|
|
|
|
if len(b.shape) == 1 or b.shape[1] == 1:
|
|
# Just one RHS
|
|
if factorize:
|
|
return self.dsolve(b)
|
|
else:
|
|
return linalg.dsolve.spsolve(self.A, b)
|
|
|
|
# Multiple RHSs
|
|
X = np.empty_like(b)
|
|
for i in range(b.shape[1]):
|
|
if factorize:
|
|
X[:,i] = self.dsolve(b[:,i])
|
|
else:
|
|
X[:,i] = linalg.dsolve.spsolve(self.A,b[:,i])
|
|
|
|
return X
|
|
|
|
def solveIter(self, b, M=None, iterSolver='CG'):
|
|
pass
|
|
|
|
def solveBackward(self, b, backend=None):
|
|
"""
|
|
Use solve instead of this interface.
|
|
|
|
Perform a backwards solve with upper triangular A in CSR format (best, if not, it will be converted).
|
|
|
|
:param str backend: which backend to use. Default is python.
|
|
:rtype: numpy.ndarray
|
|
:return: x
|
|
"""
|
|
if backend is None: backend = DEFAULTS['backward']
|
|
if type(self.A) is not sparse.csr.csr_matrix:
|
|
from scipy.sparse import csr_matrix
|
|
self.A = csr_matrix(self.A)
|
|
vals = self.A.data
|
|
rowptr = self.A.indptr
|
|
colind = self.A.indices
|
|
if backend == 'fortran':
|
|
if len(b.shape) == 1 or b.shape[1] == 1:
|
|
x = TriSolve.backward(vals, rowptr, colind, b, self.A.data.size, b.size, 1)
|
|
x = mkvc(x)
|
|
else:
|
|
x = TriSolve.backward(vals, rowptr, colind, b, self.A.data.size, b.shape[0], b.shape[1])
|
|
elif backend == 'python':
|
|
x = np.empty_like(b) # empty() is faster than zeros().
|
|
for i in reversed(xrange(self.A.shape[0])):
|
|
ith_row = vals[rowptr[i] : rowptr[i+1]]
|
|
cols = colind[rowptr[i] : rowptr[i+1]]
|
|
x_vals = x[cols]
|
|
x[i] = (b[i] - np.dot(ith_row[1:], x_vals[1:])) / ith_row[0]
|
|
return x
|
|
|
|
def solveForward(self, b, backend=None):
|
|
"""
|
|
Use solve instead of this interface.
|
|
|
|
Perform a forward solve with lower triangular A in CSR format (best, if not, it will be converted).
|
|
|
|
:param str backend: which backend to use. Default is python.
|
|
:rtype: numpy.ndarray
|
|
:return: x
|
|
"""
|
|
if backend is None: backend = DEFAULTS['forward']
|
|
if type(self.A) is not sparse.csr.csr_matrix:
|
|
from scipy.sparse import csr_matrix
|
|
self.A = csr_matrix(self.A)
|
|
vals = self.A.data
|
|
rowptr = self.A.indptr
|
|
colind = self.A.indices
|
|
if backend == 'fortran':
|
|
if len(b.shape) == 1 or b.shape[1] == 1:
|
|
x = TriSolve.forward(vals, rowptr, colind, b, self.A.data.size, b.size, 1)
|
|
x = mkvc(x)
|
|
else:
|
|
x = TriSolve.forward(vals, rowptr, colind, b, self.A.data.size, b.shape[0], b.shape[1])
|
|
elif backend == 'python':
|
|
x = np.empty_like(b) # empty() is faster than zeros().
|
|
for i in xrange(self.A.shape[0]):
|
|
ith_row = vals[rowptr[i] : rowptr[i+1]]
|
|
cols = colind[rowptr[i] : rowptr[i+1]]
|
|
x_vals = x[cols]
|
|
x[i] = (b[i] - np.dot(ith_row[:-1], x_vals[:-1])) / ith_row[-1]
|
|
return x
|
|
|
|
def solveDiagonal(self, b, backend=None):
|
|
"""
|
|
Use solve instead of this interface.
|
|
|
|
Perform a diagonal solve with diagonal matrix A.
|
|
|
|
:param str backend: which backend to use. Default is python.
|
|
:rtype: numpy.ndarray
|
|
:return: x
|
|
"""
|
|
if backend is None: backend = DEFAULTS['diagonal']
|
|
|
|
diagA = self.A.diagonal()
|
|
if len(b.shape) == 1 or b.shape[1] == 1:
|
|
# Just one RHS
|
|
return b/diagA
|
|
# Multiple RHSs
|
|
X = np.empty_like(b)
|
|
for i in range(b.shape[1]):
|
|
X[:,i] = b[:,i]/diagA
|
|
return X
|
|
|
|
|
|
if __name__ == '__main__':
|
|
from SimPEG.mesh import TensorMesh
|
|
from time import time
|
|
h1 = np.ones(20)*100.
|
|
h2 = np.ones(20)*100.
|
|
h3 = np.ones(20)*100.
|
|
|
|
h = [h1,h2,h3]
|
|
|
|
M = TensorMesh(h)
|
|
|
|
D = M.faceDiv
|
|
G = M.cellGrad
|
|
Msig = M.getFaceMass()
|
|
A = D*Msig*G
|
|
A[0,0] *= 10 # remove the constant null space from the matrix
|
|
|
|
e = np.ones(M.nC)
|
|
rhs = A.dot(e)
|
|
|
|
tic = time()
|
|
solve = Solver(A, options={'factorize':True})
|
|
x = solve.solve(rhs)
|
|
print 'Factorized', time() - tic
|
|
print np.linalg.norm(e-x,np.inf)
|
|
tic = time()
|
|
solve = Solver(A, options={'factorize':False})
|
|
x = solve.solve(rhs)
|
|
print 'spsolve', time() - tic
|
|
print np.linalg.norm(e-x,np.inf)
|
|
|
|
|
|
n = 6000
|
|
A_dense = np.random.random((n,n))
|
|
L = np.tril(np.dot(A_dense, A_dense)) # Positive definite is better conditioned.
|
|
e = np.ones(n)
|
|
b = np.dot(L, e)
|
|
|
|
A = sparse.csr_matrix(L)
|
|
pSolve = Solver(A,flag='L',options={'backend':'python'});
|
|
fSolve = Solver(A,flag='L',options={'backend':'fortran'})
|
|
tic = time()
|
|
x = pSolve.solve(b)
|
|
toc = time() - tic
|
|
print 'Error Forward Python = ', np.linalg.norm(x-e, np.inf), 'Time: ', toc
|
|
tic = time()
|
|
x = fSolve.solve(b)
|
|
toc = time() - tic
|
|
print 'Error Forward Fortran = ', np.linalg.norm(x-e, np.inf), 'Time: ', toc
|