#!/net/python/bin/python
from numpy import *
import pylab as p
from pprint import pprint
from time import clock


def delta(i,j):
    if i==j: return 1
    else: return 0

def ConstructMatrix(N):
    H = matrix(zeros([N,N]))
    for i in range(N):
        for j in range(N):
            H[i,j] = 1L * sqrt((i+1)**2+(j+1)**2)+(i+1)*delta(i,j)
    return H

def eigenvalues(H):
    return linalg.eig(H)[0]

def diagonalize(H):
    return diag(eigenvalues(H))

def randomvector(N):
    v = random.random(N)
    n = sqrt( sum(v*v) )
    return matrix(v/n).T * 1L

def lanczos(H,n,debug=False,return_v = False):
    """
    uses the Lanczos algorithm to return the first n
    eigenvalues of H
    """
    N,N2 = H.shape
    assert N == N2
    v = [matrix(zeros(N)).T,randomvector(N)]
    b = [0]
    a = [0]
    for i in range(1,n+1):
        a.append( (v[i].T*H*v[i])[0,0]) #a_i
        
        v_new = H*v[i] - b[i-1]*v[i-1] -a[i]*v[i]
        v.append (v_new/sqrt( (v_new.T*v_new)[0,0]) ) #normalized v_i+1

        b.append((v[i+1].T * H *v[i])[0,0])

        if debug:
            print '-----------'
            print H*v[i]
            print b[i-1]*v[i-1]+a[i]*v[i] + b[i]*v[i+1]

    m = diag(a[1:])
    for i in range(n-1):
        m[i,i+1] = b[i+1]
        m[i+1,i] = b[i+1]

    if return_v:
        return m,v
    else:
        return m


def print_first_N(a,N):
    try: acopy = a.copy()
    except: acopy = a[:]
    acopy.sort()
    max = minimum(N,len(a))
    for i in range(max):
        print '%1.5g\t' % acopy[i],
    print ''

def print_first_last(a):
    print '%1.9g\t' % a.min(),
    print '%1.9g' % a.max()
    

def checkconvergence(N=10,N_to_display=5):
    #checks convergence of lanczos approximation to eigenvalues
    H=ConstructMatrix(N)
    True_eigvals = eigenvalues(H)

    print 'True ',
    print_first_last(True_eigvals)
    
    h = lanczos(H,N)
    for i in range(1,N+1):
        print '%i    ' % i,
        print_first_last(eigenvalues(h[:i,:i]))

if __name__ == "__main__":
    checkconvergence(20,2)