source: trunk/anuga_work/development/2010-projects/kv/sparse.py @ 8051

Last change on this file since 8051 was 8051, checked in by steve, 13 years ago

Added in Lindon's code

File size: 10.4 KB
Line 
1"""Proof of concept sparse matrix code
2"""
3#Exactly as in the ANUGA code *except* for the initialiser for Sparse_CSR
4#Note: Sparse_CSR works *only* if there is a nonzero entry in every row!
5
6import numpy as num
7
8
9class Sparse:
10
11    def __init__(self, *args):
12        """Create sparse matrix.
13        There are two construction forms
14        Usage:
15
16        Sparse(A)     #Creates sparse matrix from dense matrix A
17        Sparse(M, N)  #Creates empty MxN sparse matrix
18        """
19
20        self.Data = {}
21           
22        if len(args) == 1:
23            try:
24                A = num.array(args[0])
25            except:
26                raise 'Input must be convertable to a numeric array'
27
28            assert len(A.shape) == 2, 'Input must be a 2d matrix'
29           
30            self.M, self.N = A.shape
31            for i in range(self.M):
32                for j in range(self.N):
33                    if A[i, j] != 0.0:
34                        self.Data[i, j] = A[i, j]
35               
36           
37        elif len(args) == 2:
38            self.M = args[0]
39            self.N = args[1]
40        else:
41            raise 'Invalid construction'
42           
43        self.shape = (self.M, self.N) 
44
45
46    def __repr__(self):
47        return '%d X %d sparse matrix:\n' %(self.M, self.N) + `self.Data`
48
49    def __len__(self):
50        """Return number of nonzeros of A
51        """
52        return len(self.Data)
53
54    def nonzeros(self):
55        """Return number of nonzeros of A
56        """       
57        return len(self)
58   
59    def __setitem__(self, key, x):
60
61        i,j = key
62        # removing these asserts will not speed things up
63        assert 0 <= i < self.M
64        assert 0 <= j < self.N       
65
66        if x != 0:
67            self.Data[key] = float(x)
68        else:
69            if self.Data.has_key( key ):           
70                del self.Data[key]
71
72    def __getitem__(self, key):
73       
74        i,j = key
75        # removing these asserts will not speed things up
76        assert 0 <= i < self.M
77        assert 0 <= j < self.N               
78
79        if self.Data.has_key( key ):
80            return self.Data[ key ]
81        else:
82            return 0.0
83
84    def copy(self):
85        #FIXME: Use the copy module instead
86        new = Sparse(self.M,self.N)
87
88        for key in self.Data.keys():
89            i, j = key
90
91            new[i,j] = self.Data[i,j]
92
93        return new
94
95
96    def todense(self):
97        D = num.zeros( (self.M, self.N), num.float)
98       
99        for i in range(self.M):
100            for j in range(self.N):
101                if self.Data.has_key( (i,j) ):               
102                    D[i, j] = self.Data[ (i,j) ]
103        return D
104
105
106   
107    def __mul__(self, other):
108        """Multiply this matrix onto 'other' which can either be
109        a numeric vector, a numeric matrix or another sparse matrix.
110        """
111
112        try:
113            B = num.array(other)
114        except:
115            msg = 'FIXME: Only numeric types implemented so far'
116            raise msg
117           
118
119        # Assume numeric types from now on
120       
121        if len(B.shape) == 0:
122            # Scalar - use __rmul__ method
123            R = B*self
124           
125        elif len(B.shape) == 1:
126            # Vector
127            msg = 'Mismatching dimensions: You cannot multiply (%d x %d) matrix onto %d-vector'\
128                  %(self.M, self.N, B.shape[0])
129            assert B.shape[0] == self.N, msg
130
131            R = num.zeros(self.M, num.float) #Result
132           
133            # Multiply nonzero elements
134            for key in self.Data.keys():
135                i, j = key
136
137                R[i] += self.Data[key]*B[j]
138        elif len(B.shape) == 2:
139       
140           
141            R = num.zeros((self.M, B.shape[1]), num.float) #Result matrix
142
143            # Multiply nonzero elements
144            for col in range(R.shape[1]):
145                # For each column
146               
147                for key in self.Data.keys():
148                    i, j = key
149
150                    R[i, col] += self.Data[key]*B[j, col]
151           
152           
153        else:
154            raise ValueError, 'Dimension too high: d=%d' %len(B.shape)
155
156        return R
157   
158
159    def __add__(self, other):
160        """Add this matrix onto 'other'
161        """
162
163        new = other.copy()
164        for key in self.Data.keys():
165            i, j = key
166
167            new[i,j] += self.Data[key]
168
169        return new
170
171
172    def __rmul__(self, other):
173        """Right multiply this matrix with scalar
174        """
175
176        try:
177            other = float(other)
178        except:
179            msg = 'Sparse matrix can only "right-multiply" onto a scalar'
180            raise TypeError, msg
181        else:
182            new = self.copy()
183            #Multiply nonzero elements
184            for key in new.Data.keys():
185                i, j = key
186
187                new.Data[key] = other*new.Data[key]
188
189        return new
190
191
192    def trans_mult(self, other):
193        """Multiply the transpose of matrix with 'other' which can be
194        a numeric vector.
195        """
196
197        try:
198            B = num.array(other)
199        except:
200            print 'FIXME: Only numeric types implemented so far'
201
202
203        #Assume numeric types from now on
204        if len(B.shape) == 1:
205            #Vector
206
207            assert B.shape[0] == self.M, 'Mismatching dimensions'
208
209            R = num.zeros((self.N,), num.float) #Result
210
211            #Multiply nonzero elements
212            for key in self.Data.keys():
213                i, j = key
214
215                R[j] += self.Data[key]*B[i]
216
217        else:
218            raise 'Can only multiply with 1d array'
219
220        return R
221
222class Sparse_CSR:
223
224    def __init__(self, A=None, data=None, Colind=None, rowptr=None, m=None, n=None):
225        """Create sparse matrix in csr format.
226
227        Sparse_CSR(A) #creates csr sparse matrix from sparse matrix
228        Matrices are not built using this format, since it's painful to
229        add values to an existing sparse_CSR instance (hence there are no
230        objects to do this.)
231
232        Rather, build a matrix, and convert it to this format for a speed
233        increase.
234
235        data - a 1D array of the data
236        Colind - The ith item in this 1D array is the column index of the
237                 ith data in the data array
238        rowptr - 1D array, with the index representing the row of the matrix.
239                 The item in the row represents the index into colind of the
240                 first data value of this row.
241                 Regard it as a pointer into the colind array, for the ith row.
242
243                 
244        """
245
246        if isinstance(A,Sparse):
247
248            keys = A.Data.keys()
249            keys.sort()
250            nnz = len(keys)
251            data    = num.zeros ( (nnz,), num.float)
252            colind  = num.zeros ( (nnz,), num.int)
253            row_ptr = num.zeros ( (A.M+1,), num.int)
254            current_row = -1
255            k = 0
256            for key in keys:
257                ikey0 = int(key[0])
258                ikey1 = int(key[1])
259                if ikey0 != current_row:
260                    current_row = ikey0
261                    row_ptr[ikey0] = k
262                data[k] = A.Data[key]
263                colind[k] = ikey1
264                k += 1
265            for row in range(current_row+1, A.M+1):
266                row_ptr[row] = nnz
267            #row_ptr[-1] = nnz
268       
269            self.data    = data
270            self.colind  = colind
271            self.row_ptr = row_ptr
272            self.M       = A.M
273            self.N       = A.N
274        elif isinstance(data,num.ndarray) and isinstance(Colind,num.ndarray) and isinstance(rowptr,num.ndarray) and isinstance(m,int) and isinstance(n,int):
275            msg = "Sparse_CSR: data is array of wrong dimensions"
276            #assert len(data.shape) == 1, msg
277            nnz = data.size
278
279            msg = "Sparse_CSR: Colind is array of wrong dimensions"
280            assert Colind.shape == (nnz,), msg
281
282            msg = "Sparse_CSR: rowptr is array of wrong dimensions"
283            assert rowptr.shape == (m+1,), msg
284
285            self.data = data
286            self.colind = Colind
287            self.row_ptr = rowptr
288            self.M = m
289            self.N = n
290        else:
291            raise ValueError, "Sparse_CSR(A) expects A == Sparse Matrix *or* data==array,colind==array,rowptr==array,m==int,n==int"
292           
293    def __repr__(self):
294        return '%d X %d sparse matrix:\n' %(self.M, self.N) + `self.data`
295
296    def __len__(self):
297        """Return number of nonzeros of A
298        """
299        return self.row_ptr[-1]
300
301    def nonzeros(self):
302        """Return number of nonzeros of A
303        """       
304        return len(self)
305
306    def todense(self):
307        D = num.zeros( (self.M, self.N), num.float)
308       
309        for i in range(self.M):
310            for ckey in range(self.row_ptr[i],self.row_ptr[i+1]):
311                j = self.colind[ckey]
312                D[i, j] = self.data[ckey]
313        return D
314
315    def __mul__(self, other):
316        """Multiply this matrix onto 'other' which can either be
317        a numeric vector, a numeric matrix or another sparse matrix.
318        """
319
320        try:
321            B = num.array(other)
322        except:
323            print 'FIXME: Only numeric types implemented so far'
324
325        return csr_mv(self,B) 
326
327
328# Setup for C extensions
329from anuga.utilities import compile
330if compile.can_use_C_extension('sparse_ext.c'):
331    # Access underlying c implementations
332    from sparse_ext import csr_mv
333
334
335if __name__ == '__main__':
336    # A little selftest
337   
338    A = Sparse(3,3)
339
340    A[1,1] = 4
341
342
343    print A
344    print A.todense()
345
346    A[1,1] = 0
347
348    print A
349    print A.todense()   
350
351    A[1,2] = 0
352
353
354    A[0,0] = 3
355    A[1,1] = 2
356    A[1,2] = 2
357    A[2,2] = 1
358
359    print A
360    print A.todense()
361
362
363    #Right hand side vector
364    v = [2,3,4]
365
366    u = A*v
367    print u
368    assert num.allclose(u, [6,14,4])
369
370    u = A.trans_mult(v)
371    print u
372    assert num.allclose(u, [6,6,10])
373
374    #Right hand side column
375    v = num.array([[2,4],[3,4],[4,4]])
376
377    u = A*v[:,0]
378    assert num.allclose(u, [6,14,4])
379
380    #u = A*v[:,1]
381    #print u
382    print A.shape
383
384    B = 3*A
385    print B.todense()
386
387    B[1,0] = 2
388
389    C = A+B
390
391    print C.todense()
392
393    C = Sparse_CSR(C)
394
395    y = C*[6,14,4]
396
397    print y
398
399    y2 = C*[[6,4],[4,28],[4,8]]
400
401    print y2
Note: See TracBrowser for help on using the repository browser.