import krypy
import numpy
import itertools
from scipy.sparse.linalg import LinearOperator
from scipy.sparse import csr_matrix

def get_spd_matrix():
    a = numpy.array(range(1,11))
    a[-1] = 1e2
    return numpy.diag(a), numpy.diag(1/a)

def get_spd_precon():
    a = numpy.array(range(1,11))
    a[-1] = 1
    return numpy.diag(a), numpy.diag(1/a)

def test_linsys_cg():
    Aorig, Ainv = get_spd_matrix()
    x = numpy.ones((10,1))
    b = numpy.dot(Aorig,x)
    Morig, Minv = get_spd_precon()
    As = [Aorig, LinearOperator(Aorig.shape, lambda x: numpy.dot(Aorig,x), dtype=numpy.double), csr_matrix(Aorig)]
    x0s = [None, numpy.zeros((10,1)), numpy.ones((10,1)), x]
    tols = [1e-14, 1e-5, 1e-2]
    maxiters = [15]
    exact_solutions = [None, x]
    for A, x0, tol, maxiter, exact_solution in itertools.product(As, x0s, tols, maxiters, exact_solutions):
        yield linsys_cg, {'A':A, 'b':b, 'x0':x0, 'tol':tol, 'maxiter': maxiter, 'exact_solution':exact_solution}

def linsys_cg(params):
    ret = krypy.linsys.cg(**params)
    assert( len(ret['relresvec'])-1 <= params['maxiter'] )
    if params['tol']>1e-15:
        assert( ret['relresvec'][-1] <= params['tol'] )

if __name__ == '__main__':
    import nose
    nose.main()

@@ -120,16 +120,344 @@
 matrix%0A%0A
+%0Adef dictproduct(d):%0A    '''enhance itertools product to process values of dicts%0A%0A    example: %0A        d = %7B'a':%5B1,2%5D,'b':%5B3,4%5D%7D%0A        then list(dictproduct(d)) == %0A        %5B%7B'a':1,'b':3%7D, %7B'a':1,'b':4%7D, %7B'a':2,'b':3%7D, %7B'a':2,'b':4%7D%5D%0A    '''%0A    for p in itertools.product(*d.values()):%0A        yield dict(zip(d.keys(), p))%0A%0A
 def get_
@@ -701,20 +701,16 @@
 ):%0A    A
-orig
 , Ainv =
@@ -773,20 +773,16 @@
 py.dot(A
-orig
 ,x)%0A    
@@ -782,20 +782,16 @@
 x)%0A    M
-orig
 , Minv =
@@ -816,20 +816,53 @@
     
-As = %5BAorig,
+params = %7B%0A        'A': %5B A, %0A               
  Lin
@@ -874,20 +874,16 @@
 erator(A
-orig
 .shape, 
@@ -907,73 +907,211 @@
 ot(A
-orig
 ,x), 
-dtype=numpy.double), csr_matrix(Aorig
+%0A                    dtype=numpy.double), %0A                csr_matrix(A)%0A                %5D,%0A        'b': %5B b, %0A                numpy.reshape(b, (b.shape%5B0%5D,)
 )%5D
+,
 %0A    
-x0s = %5BNone,
+    'x0': %5BNone, %0A               
  num
@@ -1128,16 +1128,33 @@
 10,1)), 
+%0A                
 numpy.on
@@ -1169,21 +1169,43 @@
 )), 
+%0A                
 x%5D
+,
 %0A    
-tols =
+    'tol':
  %5B1e
@@ -1224,33 +1224,44 @@
 e-2%5D
+,
 %0A    
+    '
 maxiter
-s =
+':
  %5B15%5D
-%0A
+,%0A  
     
+  '
 exac
@@ -1270,19 +1270,18 @@
 solution
-s =
+':
  %5BNone, 
@@ -1291,224 +1291,82 @@
     
-for A, x0, tol, maxiter, exact_solution in itertools.product(As, x0s, tols, maxiters, exact_solutions):%0A        yield linsys_cg, %7B'A':A, 'b':b, 'x0':x0, 'tol':tol, 'maxiter': maxiter, 'exact_solution':exact_solution%7D
+    %7D%0A    for param in dictproduct(params):%0A        yield linsys_cg, param
 %0A%0Ade

simplify tests with 'dictproduct'