"""Proof of concept sparse matrix code """ import numpy as num 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.Data = {} if len(args) == 1: try: A = num.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.Data[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.Data` def __len__(self): """Return number of nonzeros of A """ return len(self.Data) def nonzeros(self): """Return number of nonzeros of A """ return len(self) def __setitem__(self, key, x): i,j = key # removing these asserts will not speed things up assert 0 <= i < self.M assert 0 <= j < self.N if x != 0: self.Data[key] = float(x) else: if self.Data.has_key( key ): del self.Data[key] def __getitem__(self, key): i,j = key # removing these asserts will not speed things up assert 0 <= i < self.M assert 0 <= j < self.N if self.Data.has_key( key ): return self.Data[ key ] else: return 0.0 def copy(self): #FIXME: Use the copy module instead new = Sparse(self.M,self.N) for key in self.Data.keys(): i, j = key new[i,j] = self.Data[i,j] return new def todense(self): D = num.zeros( (self.M, self.N), num.float) for i in range(self.M): for j in range(self.N): if self.Data.has_key( (i,j) ): D[i, j] = self.Data[ (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. """ try: B = num.array(other) except: msg = 'FIXME: Only numeric types implemented so far' raise msg # Assume numeric types from now on if len(B.shape) == 0: # Scalar - use __rmul__ method R = B*self elif len(B.shape) == 1: # Vector msg = 'Mismatching dimensions: You cannot multiply (%d x %d) matrix onto %d-vector'\ %(self.M, self.N, B.shape[0]) assert B.shape[0] == self.N, msg R = num.zeros(self.M, num.float) #Result # Multiply nonzero elements for key in self.Data.keys(): i, j = key R[i] += self.Data[key]*B[j] elif len(B.shape) == 2: R = num.zeros((self.M, B.shape[1]), num.float) #Result matrix # Multiply nonzero elements for col in range(R.shape[1]): # For each column for key in self.Data.keys(): i, j = key R[i, col] += self.Data[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' """ new = other.copy() for key in self.Data.keys(): i, j = key new[i,j] += self.Data[key] return new def __rmul__(self, other): """Right multiply this matrix with scalar """ 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.Data.keys(): i, j = key new.Data[key] = other*new.Data[key] return new def trans_mult(self, other): """Multiply the transpose of matrix with 'other' which can be a numeric vector. """ try: B = num.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 = num.zeros((self.N,), num.float) #Result #Multiply nonzero elements for key in self.Data.keys(): i, j = key R[j] += self.Data[key]*B[i] else: raise 'Can only multiply with 1d array' return R class Sparse_CSR: def __init__(self, A): """Create sparse matrix in csr format. Sparse_CSR(A) #creates csr sparse matrix from sparse matrix Matrices are not built using this format, since it's painful to add values to an existing sparse_CSR instance (hence there are no objects to do this.) Rather, build a matrix, and convert it to this format for a speed increase. data - a 1D array of the data Colind - The ith item in this 1D array is the column index of the ith data in the data array rowptr - 1D array, with the index representing the row of the matrix. The item in the row represents the index into colind of the first data value of this row. Regard it as a pointer into the colind array, for the ith row. """ if isinstance(A,Sparse): keys = A.Data.keys() keys.sort() nnz = len(keys) data = num.zeros ( (nnz,), num.float) colind = num.zeros ( (nnz,), num.int) row_ptr = num.zeros ( (A.M+1,), num.int) current_row = -1 k = 0 for key in keys: ikey0 = int(key[0]) ikey1 = int(key[1]) if ikey0 != current_row: current_row = ikey0 row_ptr[ikey0] = k data[k] = A.Data[key] colind[k] = ikey1 k += 1 for row in range(current_row+1, A.M+1): row_ptr[row] = nnz #row_ptr[-1] = nnz self.data = data self.colind = colind self.row_ptr = row_ptr self.M = A.M self.N = A.N else: raise ValueError, "Sparse_CSR(A) expects A == Sparse Matrix" def __repr__(self): return '%d X %d sparse matrix:\n' %(self.M, self.N) + `self.data` def __len__(self): """Return number of nonzeros of A """ return self.row_ptr[-1] def nonzeros(self): """Return number of nonzeros of A """ return len(self) def todense(self): D = num.zeros( (self.M, self.N), num.float) for i in range(self.M): for ckey in range(self.row_ptr[i],self.row_ptr[i+1]): j = self.colind[ckey] D[i, j] = self.data[ckey] 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. """ try: B = num.array(other) except: print 'FIXME: Only numeric types implemented so far' return csr_mv(self,B) # Setup for C extensions from anuga.utilities import compile if compile.can_use_C_extension('sparse_ext.c'): # Access underlying c implementations from sparse_ext import csr_mv if __name__ == '__main__': # A little selftest 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 num.allclose(u, [6,14,4]) u = A.trans_mult(v) print u assert num.allclose(u, [6,6,10]) #Right hand side column v = num.array([[2,4],[3,4],[4,4]]) u = A*v[:,0] assert num.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() C = Sparse_CSR(C) y = C*[6,14,4] print y y2 = C*[[6,4],[4,28],[4,8]] print y2