Merge pull request #86 from simpeg/test_diag_est

Test diag est

SimPEG/Tests/TestUtils.py

@@ -1,7 +1,7 @@
 import numpy as np
 import matplotlib.pyplot as plt
 from numpy.linalg import norm
-from SimPEG.Utils import mkvc, sdiag
+from SimPEG.Utils import mkvc, sdiag, diagEst
 from SimPEG import Utils
 from SimPEG.Mesh import TensorMesh, LogicallyRectMesh, CylMesh
 import numpy as np

SimPEG/Tests/test_utils.py

@@ -246,5 +246,18 @@ class TestSequenceFunctions(unittest.TestCase):
         self.assertTrue(np.all(true == listArray))
         self.assertTrue(true.shape == listArray.shape)
 
+class TestDiagEst(unittest.TestCase):
+
+    def setUp(self):
+        self.n = 10
+        self.A = np.random.rand(self.n,self.n)
+        self.Adiag = np.diagonal(self.A)
+
+    def testOnes(self):
+        Adiagtest = diagEst(self.A,self.n,self.n)
+        r = np.abs(Adiagtest-self.Adiag)
+        self.assertTrue(r.dot(r) < TOL)
+
+
 if __name__ == '__main__':
     unittest.main()

SimPEG/Utils/matutils.py

@@ -341,30 +341,49 @@ def invPropertyTensor(M, tensor, returnMatrix=False):
     return T
 
 
-def diagEst(matFun, n, k=None, type='Probing'):
-    """ Based on Saad http://www-users.cs.umn.edu/~saad/PDF/umsi-2005-082.pdf, and http://www.cita.utoronto.ca/~niels/diagonal.pdf"""
+def diagEst(matFun, n, k=None, approach='Probing'):
+    """ 
+    Estimate the diagonal of a matrix, A. Note that the matrix may be a function which returns A times a vector.
+
+    Three different approaches have been implemented, 
+        1. Probing : uses cyclic permutations of vectors with ones and zeros (default)
+        2. Ones : random +/- 1 entries
+        3. Random : random vectors
+
+    :param lambda (numpy.array) matFun: matrix to estimate the diagonal of
+    :param int64 n: size of the vector that should be used to compute matFun(v)
+    :param int64 k: number of vectors to be used to estimate the diagonal
+    :param str approach: approach to be used for getting vectors
+    :rtype: numpy.array
+    :return: est_diag(A)
+
+    Based on Saad http://www-users.cs.umn.edu/~saad/PDF/umsi-2005-082.pdf, and http://www.cita.utoronto.ca/~niels/diagonal.pdf
+    """
+
+    if type(matFun).__name__=='ndarray':
+        A = matFun
+        matFun = lambda v: A.dot(v) 
 
     if k is None:
         k = np.floor(n/10.)
 
-    if type =='Ones':
+    if approach =='Ones':
         def getv(n,i=None):
             v = np.random.randn(n)
             v[v<0]  = -1.
             v[v>=0] = 1.
             return v
 
-    elif type == 'Random':
+    elif approach == 'Random':
         def getv(n,i=None):
             return np.random.randn(n)
 
-    else: #if type == 'Probing':
+    else: #if approach == 'Probing':
         def getv(n,i):
             v = np.zeros(n)
             v[i:n:k] = 1.
             return v
 
-    #d = np.zeros(n)
     Mv = np.zeros(n)
     vv = np.zeros(n)