Merge branch 'master' of https://bitbucket.org/rcockett/simpeg into Interpolation_TensorMesh

SimPEG/utils/Solver.py

@@ -0,0 +1,207 @@
+import numpy as np
+import scipy.sparse as sparse
+import scipy.sparse.linalg as linalg
+
+
+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)
+        elif self.flag == 'L':
+            return self.solveForward(b)
+        elif self.flag == 'D':
+            return self.solveDiagonal(b)
+        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='scipy'):
+        """
+            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
+        """
+        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='python'):
+        """
+            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 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
+        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='python'):
+        """
+            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 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
+        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='python'):
+        """
+            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
+        """
+        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)
+
+

SimPEG/utils/__init__.py

@@ -3,7 +3,9 @@ import sputils
 import lomutils
 import interputils
 import ModelBuilder
+import Solver
+from Solver import Solver
 from matutils import getSubArray, mkvc, ndgrid, ind2sub, sub2ind
 from sputils import spzeros, kron3, speye, sdiag
 from lomutils import volTetra, faceInfo, inv2X2BlockDiagonal, inv3X3BlockDiagonal, indexCube, exampleLomGird
-from interputils import interpmat
+from interputils import interpmat