#!/usr/bin/env python

# $Id: eigensolve.py,v 1.2 2003/11/24 15:17:14 pletzer Exp $

"""
Solve eigenvalue problem reading matrices from netCDF files
"""

import csuperlu, csparse, cvector
from Scientific.IO.NetCDF import NetCDFFile

def CCS2Dict(mat_rows, mat_cols, mat_vals, amat):
    for j in range(len(mat_cols)-1):
        for k in range(mat_cols[j], mat_cols[j+1]):
            i = mat_rows[k]
            amat[(i,j)] = mat_vals[2*k] + 1j*mat_vals[2*k+1]
    


class eigensolve:

    def __init__(self, amat_file, bmat_file):
        
        self.amat = csparse.csparse()
        self.bmat = csparse.csparse()
        
        nc = NetCDFFile(amat_file, 'r')
        mat_vals = nc.variables['values__CmPlx_Re_Im'][:]
        mat_rows = nc.variables['irow'][:]
        mat_cols = nc.variables['jcol_ptr'][:]
        CCS2Dict(mat_rows, mat_cols, mat_vals, self.amat)

        nc = NetCDFFile(bmat_file, 'r')
        mat_vals = nc.variables['values__CmPlx_Re_Im'][:]
        mat_rows = nc.variables['irow'][:]
        mat_cols = nc.variables['jcol_ptr'][:]
        CCS2Dict(mat_rows, mat_cols, mat_vals, self.bmat)

    def eigen(self, lambd0, v0=None, nmax=10, tol=1.e-10):
        if not v0:
            n = self.amat.size()[0]
            v0 = cvector.random(n, 0, 1+0j) #cvector.ones(n)
        return csuperlu.eigen(self.amat, self.bmat, lambd0, v0,
                              tol=tol, nmax=nmax, verbose=0)


#########################################################################

if __name__=='__main__':

    import sys

    lambd0 = 0.3
    nmax=10
    if len(sys.argv)>1:
        lambd0 = float(sys.argv[1])
        if len(sys.argv)>2:
            nmax = int(sys.argv[2])

    e = eigensolve('amat.nc', 'bmat.nc')
    lambd, v, residue, niter = e.eigen(lambd0, v0=None, nmax=nmax)
    print 'Lambda=%s ' % str(lambd)

    
