From 44ab7de6bcfb304fc72f6b35108452dfc53c18c0 Mon Sep 17 00:00:00 2001
From: rowanc1 <rowan@visiblegeology.com>
Date: Wed, 19 Mar 2014 10:34:37 -0700
Subject: [PATCH] Wrapper utils for Solvers

---
 SimPEG/Solver.py            | 26 +++++++++++--
 SimPEG/Utils/SolverUtils.py | 75 +++++++++++++++++++++++++++++++++++++
 SimPEG/Utils/__init__.py    |  1 +
 docs/api_Utils.rst          |  8 ++++
 4 files changed, 107 insertions(+), 3 deletions(-)
 create mode 100644 SimPEG/Utils/SolverUtils.py

diff --git a/SimPEG/Solver.py b/SimPEG/Solver.py
index 68f36798..cc6d6338 100644
--- a/SimPEG/Solver.py
+++ b/SimPEG/Solver.py
@@ -1,7 +1,7 @@
 import numpy as np
 import scipy.sparse as sp
 import scipy.sparse.linalg as linalg
-from Utils.matutils import mkvc, sdiag
+from Utils import mkvc, sdiag
 import warnings
 
 DEFAULTS = {'direct':'scipy', 'iter':'scipy', 'triangular':'fortran', 'diagonal':'python'}
@@ -214,12 +214,18 @@ class Solver(object):
     def solveIter(self, b, backend=None, M=None, iterSolver='CG', tol=1e-6, maxIter=50):
         if backend is None: backend = DEFAULTS['iter']
 
-        algorithms = {'CG':sp.linalg.cg}
+        algorithms = {'CG':sp.linalg.cg, 'QMR':sp.linalg.qmr}
         assert iterSolver in algorithms, "iterSolver must be 'CG', or implement it yourself and add it here!"
         alg = algorithms[iterSolver]
 
+        if iterSolver == 'CG':
+            opts = {'M':M}
+        elif iterSolver == 'QMR':
+            #TODO: make preconditioner better.
+            opts = {'M1':np.sqrt(M), 'M2':np.sqrt(M)}
+
         if len(b.shape) == 1 or b.shape[1] == 1:
-            x, self.info = alg(self.A, b, M=M, tol=tol, maxiter=maxIter)
+            x, self.info = alg(self.A, b,  tol=tol, maxiter=maxIter)
         else:
             x = np.empty_like(b)
             for i in range(b.shape[1]):
@@ -381,3 +387,17 @@ if __name__ == '__main__':
     toc = time() - tic
     print x
     print 'Error CG  = ', np.linalg.norm(x-e, np.inf), 'Time: ', toc, 'Info: ', iSolve.info
+
+
+
+    A = -D*D.T
+    A[0,0] *= 10 # remove the constant null space from the matrix
+    e = np.ones(M.nC)
+    b = A.dot(e)
+
+    iSolve = Solver(A, doDirect=False, options={'iterSolver': 'QMR', 'M':'J'})
+    tic = time()
+    x = iSolve.solve(b)
+    toc = time() - tic
+    print x
+    print 'Error QMR  = ', np.linalg.norm(x-e, np.inf), 'Time: ', toc, 'Info: ', iSolve.info
diff --git a/SimPEG/Utils/SolverUtils.py b/SimPEG/Utils/SolverUtils.py
new file mode 100644
index 00000000..69ca6508
--- /dev/null
+++ b/SimPEG/Utils/SolverUtils.py
@@ -0,0 +1,75 @@
+import numpy as np
+from matutils import mkvc
+import warnings
+
+def DSolverWrap(fun, factorize=True, checkAccuracy=True, accuracyTol=1e-8):
+
+    def __init__(self, A, **kwargs):
+        self.A = A.tocsc()
+        self.kwargs = kwargs
+        if factorize:
+            self.solver = fun(self.A, **kwargs)
+
+    def solve(self, b):
+        if len(b.shape) == 1 or b.shape[1] == 1:
+            # Just one RHS
+            if factorize:
+                X = self.solver.solve(b.flatten(), **self.kwargs)
+            else:
+                X = fun(self.A, b.flatten(), **self.kwargs)
+        else: # Multiple RHSs
+            X = np.empty_like(b)
+            for i in range(b.shape[1]):
+                if factorize:
+                    X[:,i] = self.solver.solve(b[:,i])
+                else:
+                    X[:,i] = fun(self.A, b[:,i], **self.kwargs)
+
+        if checkAccuracy:
+            nrm = np.linalg.norm(mkvc(self.A*X - b)) / np.linalg.norm(Utils.mkvc(b))
+            if nrm > accuracyTol:
+                msg = '### SolverWarning ###: Accuracy on solve is above tolerance: %e > %e' % (nrm, accuracyTol)
+                print msg
+                warnings.warn(msg, RuntimeWarning)
+        return X
+
+    return type(fun.__name__, (object,), {"__init__": __init__, "solve": solve})
+
+
+
+def ISolverWrap(fun, checkAccuracy=True, accuracyTol=1e-5):
+
+    def __init__(self, A, **kwargs):
+        self.A = A.tocsc()
+        self.kwargs = kwargs
+
+    def solve(self, b):
+        if len(b.shape) == 1 or b.shape[1] == 1:
+            # Just one RHS
+            out = fun(self.A, b.flatten(), **self.kwargs)
+            if type(out) is tuple and len(out) == 2:
+                # We are dealing with scipy output with an info!
+                X = out[0]
+                self.info = out[1]
+            else:
+                X = out
+        else: # Multiple RHSs
+            X = np.empty_like(b)
+            for i in range(b.shape[1]):
+                out = fun(self.A, b[:,i], **self.kwargs)
+                if type(out) is tuple and len(out) == 2:
+                    # We are dealing with scipy output with an info!
+                    X[:,i] = out[0]
+                    self.info = out[1]
+                else:
+                    X[:,i] = out
+
+        if checkAccuracy:
+            nrm = np.linalg.norm(mkvc(self.A*X - b)) / np.linalg.norm(Utils.mkvc(b))
+            if nrm > accuracyTol:
+                msg = '### SolverWarning ###: Accuracy on solve is above tolerance: %e > %e' % (nrm, accuracyTol)
+                print msg
+                warnings.warn(msg, RuntimeWarning)
+        return X
+
+    return type(fun.__name__, (object,), {"__init__": __init__, "solve": solve})
diff --git a/SimPEG/Utils/__init__.py b/SimPEG/Utils/__init__.py
index 209d1cff..56a30bc2 100644
--- a/SimPEG/Utils/__init__.py
+++ b/SimPEG/Utils/__init__.py
@@ -4,6 +4,7 @@ from lrmutils import volTetra, faceInfo, indexCube
 from interputils import interpmat
 from ipythonutils import easyAnimate as animate
 import ModelBuilder
+import SolverUtils
 
 import types
 import time
diff --git a/docs/api_Utils.rst b/docs/api_Utils.rst
index f9e24e04..0ea3e22e 100644
--- a/docs/api_Utils.rst
+++ b/docs/api_Utils.rst
@@ -15,6 +15,14 @@ Matrix Utilities
     :members:
     :undoc-members:
 
+
+Solver Utilities
+================
+
+.. automodule:: SimPEG.Utils.SolverUtils
+    :members:
+    :undoc-members:
+
 LRM Utilities
 =============