source: inundation/pyvolution/sparse.py @ 1848

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

remove method that I started to code

File size: 9.9 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
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        """
239
240        from Numeric import array, Float, Int
241
242        if isinstance(A,Sparse):
243
244            from Numeric import zeros
245            keys = A.Data.keys()
246            keys.sort()
247            nnz = len(keys)
248            data    = zeros ( (nnz,), Float)
249            colind  = zeros ( (nnz,), Int)
250            row_ptr = zeros ( (A.M+1,), Int)
251            current_row = -1
252            k = 0
253            for key in keys:
254                ikey0 = int(key[0])
255                ikey1 = int(key[1])
256                if ikey0 != current_row:
257                    current_row = ikey0
258                    row_ptr[ikey0] = k
259                data[k] = A.Data[key]
260                colind[k] = ikey1
261                k += 1
262            for row in range(current_row+1, A.M+1):
263                row_ptr[row] = nnz
264            #row_ptr[-1] = nnz
265       
266            self.data    = data
267            self.colind  = colind
268            self.row_ptr = row_ptr
269            self.M       = A.M
270            self.N       = A.N
271        else:
272            raise ValueError, "Sparse_CSR(A) expects A == Sparse Matrix"
273           
274    def __repr__(self):
275        return '%d X %d sparse matrix:\n' %(self.M, self.N) + `self.data`
276
277    def __len__(self):
278        """Return number of nonzeros of A
279        """
280        return self.row_ptr[-1]
281
282    def nonzeros(self):
283        """Return number of nonzeros of A
284        """       
285        return len(self)
286
287    def todense(self):
288        from Numeric import zeros, Float
289
290        D = zeros( (self.M, self.N), Float)
291       
292        for i in range(self.M):
293            for ckey in range(self.row_ptr[i],self.row_ptr[i+1]):
294                j = self.colind[ckey]
295                D[i, j] = self.data[ckey]
296        return D
297
298    def __mul__(self, other):
299        """Multiply this matrix onto 'other' which can either be
300        a Numeric vector, a Numeric matrix or another sparse matrix.
301        """
302
303        from Numeric import array, zeros, Float
304       
305        try:
306            B = array(other)
307        except:
308            print 'FIXME: Only Numeric types implemented so far'
309
310        return csr_mv(self,B) 
311
312
313
314def csr_mv(self, B):
315    """Python version of sparse (CSR) matrix multiplication
316    """
317
318    from Numeric import zeros, Float
319
320
321    #Assume numeric types from now on
322       
323    if len(B.shape) == 0:
324        #Scalar - use __rmul__ method
325        R = B*self
326       
327    elif len(B.shape) == 1:
328        #Vector
329        assert B.shape[0] == self.N, 'Mismatching dimensions'
330       
331        R = zeros(self.M, Float) #Result
332       
333        #Multiply nonzero elements
334        for i in range(self.M):
335            for ckey in range(self.row_ptr[i],self.row_ptr[i+1]):
336                j = self.colind[ckey]
337                R[i] += self.data[ckey]*B[j]           
338               
339    elif len(B.shape) == 2:
340       
341        R = zeros((self.M, B.shape[1]), Float) #Result matrix
342       
343        #Multiply nonzero elements
344       
345        for col in range(R.shape[1]):
346            #For each column
347            for i in range(self.M):
348                for ckey in range(self.row_ptr[i],self.row_ptr[i+1]):
349                    j = self.colind[ckey]
350                    R[i, col] += self.data[ckey]*B[j,col]           
351                   
352    else:
353        raise ValueError, 'Dimension too high: d=%d' %len(B.shape)
354   
355    return R
356
357
358
359#Setup for C extensions
360import compile
361if compile.can_use_C_extension('sparse_ext.c'):
362    #Replace python version with c implementation
363    from sparse_ext import csr_mv
364
365if __name__ == '__main__':
366
367    from Numeric import allclose, array, Float
368   
369    A = Sparse(3,3)
370
371    A[1,1] = 4
372
373
374    print A
375    print A.todense()
376
377    A[1,1] = 0
378
379    print A
380    print A.todense()   
381
382    A[1,2] = 0
383
384
385    A[0,0] = 3
386    A[1,1] = 2
387    A[1,2] = 2
388    A[2,2] = 1
389
390    print A
391    print A.todense()
392
393
394    #Right hand side vector
395    v = [2,3,4]
396
397    u = A*v
398    print u
399    assert allclose(u, [6,14,4])
400
401    u = A.trans_mult(v)
402    print u
403    assert allclose(u, [6,6,10])
404
405    #Right hand side column
406    v = array([[2,4],[3,4],[4,4]])
407
408    u = A*v[:,0]
409    assert allclose(u, [6,14,4])
410
411    #u = A*v[:,1]
412    #print u
413    print A.shape
414
415    B = 3*A
416    print B.todense()
417
418    B[1,0] = 2
419
420    C = A+B
421
422    print C.todense()
423
424    C = Sparse_CSR(C)
425
426    y = C*[6,14,4]
427
428    print y
429
430    y2 = C*[[6,4],[4,28],[4,8]]
431
432    print y2
Note: See TracBrowser for help on using the repository browser.