source: inundation/pyvolution/sparse.py @ 1916

Last change on this file since 1916 was 1849, checked in by duncan, 19 years ago

comments

File size: 10.4 KB
Line 
1"""Proof of concept sparse matrix code
2"""
3
4
5class Sparse:
6
7    def __init__(self, *args):
8        """Create sparse matrix.
9        There are two construction forms
10        Usage:
11
12        Sparse(A)     #Creates sparse matrix from dense matrix A
13        Sparse(M, N)  #Creates empty MxN sparse matrix
14
15
16       
17        """
18
19        self.Data = {}
20           
21        if len(args) == 1:
22            from Numeric import array
23            try:
24                A = 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        assert 0 <= i < self.M
63        assert 0 <= j < self.N       
64
65        if x != 0:
66            self.Data[key] = float(x)
67        else:
68            if self.Data.has_key( key ):           
69                del self.Data[key]
70
71    def __getitem__(self, key):
72       
73        i,j = key
74        assert 0 <= i < self.M
75        assert 0 <= j < self.N               
76
77        if self.Data.has_key( key ):
78            return self.Data[ key ]
79        else:
80            return 0.0
81
82    def copy(self):
83        #FIXME: Use the copy module instead
84        new = Sparse(self.M,self.N)
85
86        for key in self.Data.keys():
87            i, j = key
88
89            new[i,j] = self.Data[i,j]
90
91        return new
92
93
94    def todense(self):
95        from Numeric import zeros, Float
96
97        D = zeros( (self.M, self.N), 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        from Numeric import array, zeros, Float
113       
114        try:
115            B = array(other)
116        except:
117            msg = 'FIXME: Only Numeric types implemented so far'
118            raise msg
119           
120
121        #Assume numeric types from now on
122       
123        if len(B.shape) == 0:
124            #Scalar - use __rmul__ method
125            R = B*self
126           
127        elif len(B.shape) == 1:
128            #Vector
129            assert B.shape[0] == self.N, 'Mismatching dimensions'
130
131            R = zeros(self.M, 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 = zeros((self.M, B.shape[1]), 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        from Numeric import array, zeros, Float
164       
165        new = other.copy()
166        for key in self.Data.keys():
167            i, j = key
168
169            new[i,j] += self.Data[key]
170
171        return new
172
173
174    def __rmul__(self, other):
175        """Right multiply this matrix with scalar
176        """
177
178        from Numeric import array, zeros, Float
179       
180        try:
181            other = float(other)
182        except:
183            msg = 'Sparse matrix can only "right-multiply" onto a scalar'
184            raise TypeError, msg
185        else:
186            new = self.copy()
187            #Multiply nonzero elements
188            for key in new.Data.keys():
189                i, j = key
190
191                new.Data[key] = other*new.Data[key]
192
193        return new
194
195
196    def trans_mult(self, other):
197        """Multiply the transpose of matrix with 'other' which can be
198        a Numeric vector.
199        """
200
201        from Numeric import array, zeros, Float
202       
203        try:
204            B = array(other)
205        except:
206            print 'FIXME: Only Numeric types implemented so far'
207
208
209        #Assume numeric types from now on
210        if len(B.shape) == 1:
211            #Vector
212
213            assert B.shape[0] == self.M, 'Mismatching dimensions'
214
215            R = zeros((self.N,), Float) #Result
216
217            #Multiply nonzero elements
218            for key in self.Data.keys():
219                i, j = key
220
221                R[j] += self.Data[key]*B[i]
222
223        else:
224            raise 'Can only multiply with 1d array'
225
226        return R
227
228class Sparse_CSR:
229
230    def __init__(self, A):
231        """Create sparse matrix in csr format.
232
233        Sparse_CSR(A) #creates csr sparse matrix from sparse matrix
234        Matrices are not built using this format, since it's painful to
235        add values to an existing sparse_CSR instance (hence there are no
236        objects to do this.)
237
238        Rather, build a matrix, and convert it to this format for a speed
239        increase.
240
241        data - a 1D array of the data
242        Colind - The ith item in this 1D array is the column index of the
243                 ith data in the data array
244        rowptr - 1D array, with the index representing the row of the matrix.
245                 The item in the row represents the index into colind of the
246                 first data value of this row.
247                 Regard it as a pointer into the colind array, for the ith row.
248
249                 
250        """
251
252        from Numeric import array, Float, Int
253
254        if isinstance(A,Sparse):
255
256            from Numeric import zeros
257            keys = A.Data.keys()
258            keys.sort()
259            nnz = len(keys)
260            data    = zeros ( (nnz,), Float)
261            colind  = zeros ( (nnz,), Int)
262            row_ptr = zeros ( (A.M+1,), Int)
263            current_row = -1
264            k = 0
265            for key in keys:
266                ikey0 = int(key[0])
267                ikey1 = int(key[1])
268                if ikey0 != current_row:
269                    current_row = ikey0
270                    row_ptr[ikey0] = k
271                data[k] = A.Data[key]
272                colind[k] = ikey1
273                k += 1
274            for row in range(current_row+1, A.M+1):
275                row_ptr[row] = nnz
276            #row_ptr[-1] = nnz
277       
278            self.data    = data
279            self.colind  = colind
280            self.row_ptr = row_ptr
281            self.M       = A.M
282            self.N       = A.N
283        else:
284            raise ValueError, "Sparse_CSR(A) expects A == Sparse Matrix"
285           
286    def __repr__(self):
287        return '%d X %d sparse matrix:\n' %(self.M, self.N) + `self.data`
288
289    def __len__(self):
290        """Return number of nonzeros of A
291        """
292        return self.row_ptr[-1]
293
294    def nonzeros(self):
295        """Return number of nonzeros of A
296        """       
297        return len(self)
298
299    def todense(self):
300        from Numeric import zeros, Float
301
302        D = zeros( (self.M, self.N), Float)
303       
304        for i in range(self.M):
305            for ckey in range(self.row_ptr[i],self.row_ptr[i+1]):
306                j = self.colind[ckey]
307                D[i, j] = self.data[ckey]
308        return D
309
310    def __mul__(self, other):
311        """Multiply this matrix onto 'other' which can either be
312        a Numeric vector, a Numeric matrix or another sparse matrix.
313        """
314
315        from Numeric import array, zeros, Float
316       
317        try:
318            B = array(other)
319        except:
320            print 'FIXME: Only Numeric types implemented so far'
321
322        return csr_mv(self,B) 
323
324
325
326def csr_mv(self, B):
327    """Python version of sparse (CSR) matrix multiplication
328    """
329
330    from Numeric import zeros, Float
331
332
333    #Assume numeric types from now on
334       
335    if len(B.shape) == 0:
336        #Scalar - use __rmul__ method
337        R = B*self
338       
339    elif len(B.shape) == 1:
340        #Vector
341        assert B.shape[0] == self.N, 'Mismatching dimensions'
342       
343        R = zeros(self.M, Float) #Result
344       
345        #Multiply nonzero elements
346        for i in range(self.M):
347            for ckey in range(self.row_ptr[i],self.row_ptr[i+1]):
348                j = self.colind[ckey]
349                R[i] += self.data[ckey]*B[j]           
350               
351    elif len(B.shape) == 2:
352       
353        R = zeros((self.M, B.shape[1]), Float) #Result matrix
354       
355        #Multiply nonzero elements
356       
357        for col in range(R.shape[1]):
358            #For each column
359            for i in range(self.M):
360                for ckey in range(self.row_ptr[i],self.row_ptr[i+1]):
361                    j = self.colind[ckey]
362                    R[i, col] += self.data[ckey]*B[j,col]           
363                   
364    else:
365        raise ValueError, 'Dimension too high: d=%d' %len(B.shape)
366   
367    return R
368
369
370
371#Setup for C extensions
372import compile
373if compile.can_use_C_extension('sparse_ext.c'):
374    #Replace python version with c implementation
375    from sparse_ext import csr_mv
376
377if __name__ == '__main__':
378
379    from Numeric import allclose, array, Float
380   
381    A = Sparse(3,3)
382
383    A[1,1] = 4
384
385
386    print A
387    print A.todense()
388
389    A[1,1] = 0
390
391    print A
392    print A.todense()   
393
394    A[1,2] = 0
395
396
397    A[0,0] = 3
398    A[1,1] = 2
399    A[1,2] = 2
400    A[2,2] = 1
401
402    print A
403    print A.todense()
404
405
406    #Right hand side vector
407    v = [2,3,4]
408
409    u = A*v
410    print u
411    assert allclose(u, [6,14,4])
412
413    u = A.trans_mult(v)
414    print u
415    assert allclose(u, [6,6,10])
416
417    #Right hand side column
418    v = array([[2,4],[3,4],[4,4]])
419
420    u = A*v[:,0]
421    assert allclose(u, [6,14,4])
422
423    #u = A*v[:,1]
424    #print u
425    print A.shape
426
427    B = 3*A
428    print B.todense()
429
430    B[1,0] = 2
431
432    C = A+B
433
434    print C.todense()
435
436    C = Sparse_CSR(C)
437
438    y = C*[6,14,4]
439
440    print y
441
442    y2 = C*[[6,4],[4,28],[4,8]]
443
444    print y2
Note: See TracBrowser for help on using the repository browser.