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

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


class Sparse:

    def __init__(self, *args):
        """Create sparse matrix.
        There are two construction forms
        Usage:

        Sparse(A)     #Creates sparse matrix from dense matrix A 
        Sparse(M, N)  #Creates empty MxN sparse matrix


        
        """

        self.A = {}
            
        if len(args) == 1:
            from Numeric import array
            try:
                A = array(args[0])
            except:
                raise 'Input must be convertable to a Numeric array'

            assert len(A.shape) == 2, 'Input must be a 2d matrix'
            
            self.M, self.N = A.shape
            for i in range(self.M):
                for j in range(self.N):
                    if A[i, j] != 0.0:
                        self.A[i, j] = A[i, j]
                
            
        elif len(args) == 2:
            self.M = args[0]
            self.N = args[1]
        else:
            raise 'Invalid construction'
            
        self.shape = (self.M, self.N) 


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

    def __len__(self):
        """Return number of nonzeros of A
        """
        return len(self.A)

    def nonzeros(self):
        """Return number of nonzeros of A
        """        
        return len(self)
    
    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):
        #FIXME: Use the copy module instead
        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
        
        if len(B.shape) == 0:
            #Scalar - use __rmul__ method
            R = B*self
            
        elif len(B.shape) == 1:
            #Vector
            assert B.shape[0] == self.N, 'Mismatching dimensions'

            R = zeros(self.M, Float) #Result
            
            #Multiply nonzero elements
            for key in self.A.keys():
                i, j = key

                R[i] += self.A[key]*B[j]
        elif len(B.shape) == 2:
        
            
            R = zeros((self.M, B.shape[1]), Float) #Result matrix

            #Multiply nonzero elements
            for col in range(R.shape[1]):
                #For each column
                
                for key in self.A.keys():
                    i, j = key

                    R[i, col] += self.A[key]*B[j, col]
            
            
        else:
            raise ValueError, 'Dimension too high: d=%d' %len(B.shape)

        return R
    

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

        from Numeric import array, zeros, Float
        
        new = other.copy()
        for key in self.A.keys():
            i, j = key

            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:
            msg = 'Sparse matrix can only "right-multiply" onto a scalar'
            raise TypeError, msg
        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()