source: inundation/ga/storm_surge/pyvolution/sparse.py @ 471

"""Proof of concept sparse matrix code
"""

#from scipy_base import * #Hardly worth importing scipy just to get isscalar.
from cg_solve import conjugate_gradient, VectorShapeError

class Sparse:

    def __init__(self, M, N):
        """Set dimensions
        """

        self.M = M
        self.N = N
        self.shape = (M,N)
        self.A = {}

    def __repr__(self):
        return '%d X %d sparse matrix:\n' %(self.M, self.N) + `self.A` 

    def __setitem__(self, key, x):

        i,j = key
        assert 0 <= i < self.M
        assert 0 <= j < self.N        

        if x != 0:
            self.A[key] = float(x)
        else:
            if self.A.has_key( key ):            
                del self.A[key]

    def __getitem__(self, key):
        
        i,j = key
        assert 0 <= i < self.M
        assert 0 <= j < self.N                

        if self.A.has_key( key ):
            return self.A[ key ]
        else:
            return 0.0

    def copy(self):

        new = Sparse(self.M,self.N)

        for key in self.A.keys():
            i, j = key

            new[i,j] = self.A[i,j]

        return new


    def todense(self):
        from Numeric import zeros, Float

        D = zeros( (self.M, self.N), Float)
        
        for i in range(self.M):
            for j in range(self.N):
                if self.A.has_key( (i,j) ):                
                    D[i, j] = self.A[ (i,j) ]
        return D


    def __mul__(self, other):
        """Multiply this matrix onto 'other' which can either be
        a Numeric vector, a Numeric matrix or another sparse matrix.
        """

        from Numeric import array, zeros, Float
        
        try:
            B = array(other)
        except:
            print 'FIXME: Only Numeric types implemented so far'


        #Assume numeric types from now on
        R = zeros((self.M,), Float) #Result
        
        if len(B.shape) == 1:
            #Vector

##             print 'B.shape      ',B.shape
##             print 'self.shape ',self.shape
            
            assert B.shape[0] == self.N, 'Mismatching dimensions'

            #Multiply nonzero elements
            for key in self.A.keys():
                i, j = key

                R[i] += self.A[key]*B[j]

        else:
            raise ValueError, 'Numeric matrix not yet implemented'

        return R

    def __add__(self, other):
        """Add this matrix onto 'other' 
        """

        from Numeric import array, zeros, Float
        
        new = other.copy()

 #       print 'self.shape',self.shape
 #       print 'other.shape',other.shape
        
        for key in self.A.keys():
            i, j = key

            new[i,j] = new[i,j] + self.A[key]

        return new


    def __rmul__(self, other):
        """Right multiply this matrix with scalar
        """

        from Numeric import array, zeros, Float
        
        try:
            other = float(other)
        except:
            raise 'only right multiple with scalar implemented'
        else:
            new = self.copy()
            #Multiply nonzero elements
            for key in new.A.keys():
                i, j = key

                new.A[key] = other*new.A[key]

        return new


    def trans_mult(self, other):
        """Multiply the transpose of  matrix with 'other' which can be
        a Numeric vector.
        """

        from Numeric import array, zeros, Float
        
        try:
            B = array(other)
        except:
            print 'FIXME: Only Numeric types implemented so far'


        #Assume numeric types from now on

        
        if len(B.shape) == 1:
            #Vector

            assert B.shape[0] == self.M, 'Mismatching dimensions'

            R = zeros((self.N,), Float) #Result

            #Multiply nonzero elements
            for key in self.A.keys():
                i, j = key

                R[j] += self.A[key]*B[i]

        else:
            raise 'Can only multiply with 1d array'

        return R


    
    
if __name__ == '__main__':

    from Numeric import allclose, array, Float 
    
    A = Sparse(3,3)

    A[1,1] = 4


    print A
    print A.todense()

    A[1,1] = 0

    print A
    print A.todense()    

    A[1,2] = 0


    A[0,0] = 3
    A[1,1] = 2
    A[1,2] = 2
    A[2,2] = 1

    print A
    print A.todense()


    #Right hand side vector
    v = [2,3,4]

    u = A*v
    print u
    assert allclose(u, [6,14,4])

    u = A.trans_mult(v)
    print u
    assert allclose(u, [6,6,10])

    #Right hand side column
    v = array([[2,4],[3,4],[4,4]])

    u = A*v[:,0]
    assert allclose(u, [6,14,4])

    #u = A*v[:,1]
    #print u
    print A.shape

    B = 3*A
    print B.todense()

    B[1,0] = 2

    C = A+B

    print C.todense()