from scipy import *
from scipy.sandbox import lobpcg
from symeig import symeig
from pylab import plot, show, legend, xlabel, ylabel
set_printoptions(precision=3,linewidth=90)
import time

def test(n):
  x = arange(1,n+1)
  B = diag(1./x)
  y = arange(n-1,0,-1)
  z = arange(2*n-1,0,-2)
  A = diag(z)-diag(y,-1)-diag(y,1)
  return A,B

def as2d( ar ):
    if ar.ndim == 2:
        return ar
    else: # Assume 1!
        aux = nm.array( ar, copy = False )
        aux.shape = (ar.shape[0], 1)
        return aux

def precond(x):
    y= linalg.cho_solve((LorU, lower),x)
    return as2d(y)

m = 10  # Blocksize
N = array(([128,256,512,1024,2048])) # Increasing matrix size

data1=[]
data2=[]

for n in N:
  print '******', n
  A,B = test(n) # Mikota pair 
  X = rand(n,m)
  X = linalg.orth(X)

  tt = time.clock()
  (LorU, lower) = linalg.cho_factor(A, lower=0, overwrite_a=0)
  eigs,vecs = lobpcg.lobpcg(X,A,B,operatorT = precond,
                            residualTolerance = 1e-4, maxIterations = 40)
  data1.append(time.clock()-tt)
  eigs = sort(eigs)
  print 
  print 'Results by LOBPCG'
  print
  print n,eigs

  tt = time.clock()
  w,v=symeig(A,B,range=(1,m))
  data2.append(time.clock()-tt)
  print 
  print 'Results by symeig'
  print 
  print n, w

xlabel(r'Size $n$')
ylabel(r'Elapsed time $t$')
plot(N,data1,label='LOBPCG')
plot(N,data2,label='SYMEIG')
legend()
show()