xref: /petsc/src/binding/petsc4py/test/test_mat_py.py (revision 81b34b0469c4c918ac4feb79d232c68b446ecfbd) 
1from petsc4py import PETSc
2import unittest, numpy
3from sys import getrefcount
4# --------------------------------------------------------------------
5
6class Matrix(object):
7
8    def __init__(self):
9        pass
10
11    def create(self, mat):
12        pass
13
14    def destroy(self, mat):
15        pass
16
17class ScaledIdentity(Matrix):
18
19    s = 2.0
20
21    def scale(self, mat, s):
22        self.s *= s
23
24    def shift(self, mat, s):
25        self.s += s
26
27    def mult(self, mat, x, y):
28        x.copy(y)
29        y.scale(self.s)
30
31    def getDiagonal(self, mat, vd):
32        vd.set(self.s)
33
34    def productSetFromOptions(self, mat, producttype, A, B, C):
35        return True
36
37    def productSymbolic(self, mat, product, producttype, A, B, C):
38        if producttype == 'AB':
39            if mat is A: # product = identity * B
40                product.setType(B.getType())
41                product.setSizes(B.getSizes())
42                product.setUp()
43                product.assemble()
44                B.copy(product)
45            elif mat is B: # product = A * identity
46                product.setType(A.getType())
47                product.setSizes(A.getSizes())
48                product.setUp()
49                product.assemble()
50                A.copy(product)
51            else:
52                raise RuntimeError('wrong configuration')
53        elif producttype == 'AtB':
54            if mat is A: # product = identity^T * B
55                product.setType(B.getType())
56                product.setSizes(B.getSizes())
57                product.setUp()
58                product.assemble()
59                B.copy(product)
60            elif mat is B: # product = A^T * identity
61                tmp = PETSc.Mat()
62                A.transpose(tmp)
63                product.setType(tmp.getType())
64                product.setSizes(tmp.getSizes())
65                product.setUp()
66                product.assemble()
67                tmp.copy(product)
68            else:
69                raise RuntimeError('wrong configuration')
70        elif producttype == 'ABt':
71            if mat is A: # product = identity * B^T
72                tmp = PETSc.Mat()
73                B.transpose(tmp)
74                product.setType(tmp.getType())
75                product.setSizes(tmp.getSizes())
76                product.setUp()
77                product.assemble()
78                tmp.copy(product)
79            elif mat is B: # product = A * identity^T
80                product.setType(A.getType())
81                product.setSizes(A.getSizes())
82                product.setUp()
83                product.assemble()
84                A.copy(product)
85            else:
86                raise RuntimeError('wrong configuration')
87        elif producttype == 'PtAP':
88            if mat is A: # product = P^T * identity * P
89                self.tmp = PETSc.Mat()
90                B.transposeMatMult(B, self.tmp)
91                product.setType(self.tmp.getType())
92                product.setSizes(self.tmp.getSizes())
93                product.setUp()
94                product.assemble()
95                self.tmp.copy(product)
96            elif mat is B: # product = identity^T * A * identity
97                product.setType(A.getType())
98                product.setSizes(A.getSizes())
99                product.setUp()
100                product.assemble()
101                A.copy(product)
102            else:
103                raise RuntimeError('wrong configuration')
104        elif producttype == 'RARt':
105            if mat is A: # product = R * identity * R^t
106                self.tmp = PETSc.Mat()
107                B.matTransposeMult(B, self.tmp)
108                product.setType(self.tmp.getType())
109                product.setSizes(self.tmp.getSizes())
110                product.setUp()
111                product.assemble()
112                self.tmp.copy(product)
113            elif mat is B: # product = identity * A * identity^T
114                product.setType(A.getType())
115                product.setSizes(A.getSizes())
116                product.setUp()
117                product.assemble()
118                A.copy(product)
119            else:
120                raise RuntimeError('wrong configuration')
121        elif producttype == 'ABC':
122            if mat is A: # product = identity * B * C
123                self.tmp = PETSc.Mat()
124                B.matMult(C, self.tmp)
125                product.setType(self.tmp.getType())
126                product.setSizes(self.tmp.getSizes())
127                product.setUp()
128                product.assemble()
129                self.tmp.copy(product)
130            elif mat is B: # product = A * identity * C
131                self.tmp = PETSc.Mat()
132                A.matMult(C, self.tmp)
133                product.setType(self.tmp.getType())
134                product.setSizes(self.tmp.getSizes())
135                product.setUp()
136                product.assemble()
137                self.tmp.copy(product)
138            elif mat is C: # product = A * B * identity
139                self.tmp = PETSc.Mat()
140                A.matMult(B, self.tmp)
141                product.setType(self.tmp.getType())
142                product.setSizes(self.tmp.getSizes())
143                product.setUp()
144                product.assemble()
145                self.tmp.copy(product)
146            else:
147                raise RuntimeError('wrong configuration')
148        else:
149            raise RuntimeError('Product {} not implemented'.format(producttype))
150        product.zeroEntries()
151
152    def productNumeric(self, mat, product, producttype, A, B, C):
153        if producttype == 'AB':
154            if mat is A: # product = identity * B
155                B.copy(product, structure=True)
156            elif mat is B: # product = A * identity
157                A.copy(product, structure=True)
158            else:
159                raise RuntimeError('wrong configuration')
160            product.scale(self.s)
161        elif producttype == 'AtB':
162            if mat is A: # product = identity^T * B
163                B.copy(product, structure=True)
164            elif mat is B: # product = A^T * identity
165                A.transpose(product)
166            else:
167                raise RuntimeError('wrong configuration')
168            product.scale(self.s)
169        elif producttype == 'ABt':
170            if mat is A: # product = identity * B^T
171                B.transpose(product)
172            elif mat is B: # product = A * identity^T
173                A.copy(product, structure=True)
174            else:
175                raise RuntimeError('wrong configuration')
176            product.scale(self.s)
177        elif producttype == 'PtAP':
178            if mat is A: # product = P^T * identity * P
179                B.transposeMatMult(B, self.tmp)
180                self.tmp.copy(product, structure=True)
181                product.scale(self.s)
182            elif mat is B: # product = identity^T * A * identity
183                A.copy(product, structure=True)
184                product.scale(self.s**2)
185            else:
186                raise RuntimeError('wrong configuration')
187        elif producttype == 'RARt':
188            if mat is A: # product = R * identity * R^t
189                B.matTransposeMult(B, self.tmp)
190                self.tmp.copy(product, structure=True)
191                product.scale(self.s)
192            elif mat is B: # product = identity * A * identity^T
193                A.copy(product, structure=True)
194                product.scale(self.s**2)
195            else:
196                raise RuntimeError('wrong configuration')
197        elif producttype == 'ABC':
198            if mat is A: # product = identity * B * C
199                B.matMult(C, self.tmp)
200                self.tmp.copy(product, structure=True)
201            elif mat is B: # product = A * identity * C
202                A.matMult(C, self.tmp)
203                self.tmp.copy(product, structure=True)
204            elif mat is C: # product = A * B * identity
205                A.matMult(B, self.tmp)
206                self.tmp.copy(product, structure=True)
207            else:
208                raise RuntimeError('wrong configuration')
209            product.scale(self.s)
210        else:
211            raise RuntimeError('Product {} not implemented'.format(producttype))
212
213class Diagonal(Matrix):
214
215    def create(self, mat):
216        super(Diagonal,self).create(mat)
217        mat.setUp()
218        self.D = mat.createVecLeft()
219
220    def destroy(self, mat):
221        self.D.destroy()
222        super(Diagonal,self).destroy(mat)
223
224    def scale(self, mat, a):
225        self.D.scale(a)
226
227    def shift(self, mat, a):
228        self.D.shift(a)
229
230    def zeroEntries(self, mat):
231        self.D.zeroEntries()
232
233    def mult(self, mat, x, y):
234        y.pointwiseMult(x, self.D)
235
236    def getDiagonal(self, mat, vd):
237        self.D.copy(vd)
238
239    def setDiagonal(self, mat, vd, im):
240        if isinstance (im, bool):
241            addv = im
242            if addv:
243                self.D.axpy(1, vd)
244            else:
245                vd.copy(self.D)
246        elif im == PETSc.InsertMode.INSERT_VALUES:
247            vd.copy(self.D)
248        elif im == PETSc.InsertMode.ADD_VALUES:
249            self.D.axpy(1, vd)
250        else:
251            raise ValueError('wrong InsertMode %d'% im)
252
253    def diagonalScale(self, mat, vl, vr):
254        if vl: self.D.pointwiseMult(self.D, vl)
255        if vr: self.D.pointwiseMult(self.D, vr)
256
257# --------------------------------------------------------------------
258
259class TestMatrix(unittest.TestCase):
260
261    COMM = PETSc.COMM_WORLD
262    PYMOD = __name__
263    PYCLS = 'Matrix'
264
265    def _getCtx(self):
266        return self.A.getPythonContext()
267
268    def setUp(self):
269        N = self.N = 10
270        self.A = PETSc.Mat()
271        if 0: # command line way
272            self.A.create(self.COMM)
273            self.A.setSizes([N,N])
274            self.A.setType('python')
275            OptDB = PETSc.Options(self.A)
276            OptDB['mat_python_type'] = '%s.%s' % (self.PYMOD,self.PYCLS)
277            self.A.setFromOptions()
278            self.A.setUp()
279            del OptDB['mat_python_type']
280            self.assertTrue(self._getCtx() is not None)
281        else: # python way
282            context = globals()[self.PYCLS]()
283            self.A.createPython([N,N], context, comm=self.COMM)
284            self.A.setUp()
285            self.assertTrue(self._getCtx() is context)
286            self.assertEqual(getrefcount(context), 3)
287            del context
288            self.assertEqual(getrefcount(self._getCtx()), 2)
289
290    def tearDown(self):
291        ctx = self.A.getPythonContext()
292        self.assertEqual(getrefcount(ctx), 3)
293        self.A.destroy() # XXX
294        self.A = None
295        self.assertEqual(getrefcount(ctx), 2)
296        #import gc,pprint; pprint.pprint(gc.get_referrers(ctx))
297
298    def testBasic(self):
299        ctx = self.A.getPythonContext()
300        self.assertTrue(self._getCtx() is ctx)
301        self.assertEqual(getrefcount(ctx), 3)
302
303    def testZeroEntries(self):
304        f = lambda : self.A.zeroEntries()
305        self.assertRaises(Exception, f)
306
307    def testMult(self):
308        x, y = self.A.createVecs()
309        f = lambda : self.A.mult(x, y)
310        self.assertRaises(Exception, f)
311
312    def testMultTranspose(self):
313        x, y = self.A.createVecs()
314        f = lambda : self.A.multTranspose(x, y)
315        self.assertRaises(Exception, f)
316
317    def testGetDiagonal(self):
318        d = self.A.createVecLeft()
319        f = lambda : self.A.getDiagonal(d)
320        self.assertRaises(Exception, f)
321
322    def testSetDiagonal(self):
323        d = self.A.createVecLeft()
324        f = lambda : self.A.setDiagonal(d)
325        self.assertRaises(Exception, f)
326
327    def testDiagonalScale(self):
328        x, y = self.A.createVecs()
329        f = lambda : self.A.diagonalScale(x, y)
330        self.assertRaises(Exception, f)
331
332    def testSetVecType(self):
333        self.A.setVecType('mpi')
334        self.assertTrue('mpi' == self.A.getVecType())
335
336class TestScaledIdentity(TestMatrix):
337
338    PYCLS = 'ScaledIdentity'
339
340    def testMult(self):
341        s = self._getCtx().s
342        x, y = self.A.createVecs()
343        x.setRandom()
344        self.A.mult(x,y)
345        self.assertTrue(y.equal(s*x))
346
347    def testMultTransposeSymmKnown(self):
348        s = self._getCtx().s
349        x, y = self.A.createVecs()
350        x.setRandom()
351        self.A.setOption(PETSc.Mat.Option.SYMMETRIC, True)
352        self.A.multTranspose(x,y)
353        self.assertTrue(y.equal(s*x))
354        self.A.setOption(PETSc.Mat.Option.SYMMETRIC, False)
355        f = lambda : self.A.multTranspose(x, y)
356        self.assertRaises(Exception, f)
357
358    def testMultTransposeNewMeth(self):
359        s = self._getCtx().s
360        x, y = self.A.createVecs()
361        x.setRandom()
362        AA = self.A.getPythonContext()
363        AA.multTranspose = AA.mult
364        self.A.multTranspose(x,y)
365        del AA.multTranspose
366        self.assertTrue(y.equal(s*x))
367
368    def testGetDiagonal(self):
369        s = self._getCtx().s
370        d = self.A.createVecLeft()
371        o = d.duplicate()
372        o.set(s)
373        self.A.getDiagonal(d)
374        self.assertTrue(o.equal(d))
375
376    def testMatMat(self):
377        s = self._getCtx().s
378        R = PETSc.Random().create(self.COMM)
379        R.setFromOptions()
380        A = PETSc.Mat().create(self.COMM)
381        A.setSizes(self.A.getSizes())
382        A.setType(PETSc.Mat.Type.AIJ)
383        A.setUp()
384        A.setRandom(R)
385        B = PETSc.Mat().create(self.COMM)
386        B.setSizes(self.A.getSizes())
387        B.setType(PETSc.Mat.Type.AIJ)
388        B.setUp()
389        B.setRandom(R)
390        I = PETSc.Mat().create(self.COMM)
391        I.setSizes(self.A.getSizes())
392        I.setType(PETSc.Mat.Type.AIJ)
393        I.setUp()
394        I.assemble()
395        I.shift(s)
396
397        self.assertTrue(self.A.matMult(A).equal(I.matMult(A)))
398        self.assertTrue(A.matMult(self.A).equal(A.matMult(I)))
399        if self.A.getComm().Get_size() == 1:
400            self.assertTrue(self.A.matTransposeMult(A).equal(I.matTransposeMult(A)))
401            self.assertTrue(A.matTransposeMult(self.A).equal(A.matTransposeMult(I)))
402        self.assertTrue(self.A.transposeMatMult(A).equal(I.transposeMatMult(A)))
403        self.assertTrue(A.transposeMatMult(self.A).equal(A.transposeMatMult(I)))
404        self.assertAlmostEqual((self.A.ptap(A) - I.ptap(A)).norm(), 0.0, places=5)
405        self.assertAlmostEqual((A.ptap(self.A) - A.ptap(I)).norm(), 0.0, places=5)
406        if self.A.getComm().Get_size() == 1:
407            self.assertAlmostEqual((self.A.rart(A) - I.rart(A)).norm(), 0.0, places=5)
408            self.assertAlmostEqual((A.rart(self.A) - A.rart(I)).norm(), 0.0, places=5)
409        self.assertAlmostEqual((self.A.matMatMult(A,B)-I.matMatMult(A,B)).norm(), 0.0, places=5)
410        self.assertAlmostEqual((A.matMatMult(self.A,B)-A.matMatMult(I,B)).norm(), 0.0, places=5)
411        self.assertAlmostEqual((A.matMatMult(B,self.A)-A.matMatMult(B,I)).norm(), 0.0, places=5)
412
413    def testH2Opus(self):
414        if not PETSc.Sys.hasExternalPackage("h2opus"):
415            return
416        if self.A.getComm().Get_size() > 1:
417            return
418        h = PETSc.Mat()
419
420        # need transpose operation for norm estimation
421        AA = self.A.getPythonContext()
422        AA.multTranspose = AA.mult
423
424        # without coordinates
425        h.createH2OpusFromMat(self.A,leafsize=2)
426        h.assemble()
427        h.destroy()
428
429        # with coordinates
430        coords = numpy.linspace((1,2,3),(10,20,30),self.A.getSize()[0],dtype=PETSc.RealType)
431        h.createH2OpusFromMat(self.A,coords,leafsize=2)
432        h.assemble()
433        h.destroy()
434
435        del AA.multTranspose
436
437    def testShift(self):
438        sold = self._getCtx().s
439        self.A.shift(-0.5)
440        s = self._getCtx().s
441        self.assertTrue(s == sold - 0.5)
442
443    def testScale(self):
444        sold = self._getCtx().s
445        self.A.scale(-0.5)
446        s = self._getCtx().s
447        self.assertTrue(s == sold * -0.5)
448
449class TestDiagonal(TestMatrix):
450
451    PYCLS = 'Diagonal'
452
453    def setUp(self):
454        super(TestDiagonal, self).setUp()
455        D = self.A.createVecLeft()
456        s, e = D.getOwnershipRange()
457        for i in range(s, e):
458            D[i] = i+1
459        D.assemble()
460        self.A.setDiagonal(D)
461
462    def testZeroEntries(self):
463        self.A.zeroEntries()
464        D = self._getCtx().D
465        self.assertEqual(D.norm(), 0)
466
467    def testMult(self):
468        x, y = self.A.createVecs()
469        x.set(1)
470        self.A.mult(x,y)
471        self.assertTrue(y.equal(self._getCtx().D))
472
473    def testMultTransposeSymmKnown(self):
474        x, y = self.A.createVecs()
475        x.set(1)
476        self.A.setOption(PETSc.Mat.Option.SYMMETRIC, True)
477        self.A.multTranspose(x,y)
478        self.assertTrue(y.equal(self._getCtx().D))
479        self.A.setOption(PETSc.Mat.Option.SYMMETRIC, False)
480        f = lambda : self.A.multTranspose(x, y)
481        self.assertRaises(Exception, f)
482
483    def testMultTransposeNewMeth(self):
484        x, y = self.A.createVecs()
485        x.set(1)
486        AA = self.A.getPythonContext()
487        AA.multTranspose = AA.mult
488        self.A.multTranspose(x,y)
489        del AA.multTranspose
490        self.assertTrue(y.equal(self._getCtx().D))
491
492    def testGetDiagonal(self):
493        d = self.A.createVecLeft()
494        self.A.getDiagonal(d)
495        self.assertTrue(d.equal(self._getCtx().D))
496
497    def testSetDiagonal(self):
498        d = self.A.createVecLeft()
499        d.setRandom()
500        self.A.setDiagonal(d)
501        self.assertTrue(d.equal(self._getCtx().D))
502
503    def testDiagonalScale(self):
504        x, y = self.A.createVecs()
505        x.set(2)
506        y.set(3)
507        old = self._getCtx().D.copy()
508        self.A.diagonalScale(x, y)
509        D = self._getCtx().D
510        self.assertTrue(D.equal(old*6))
511
512    def testCreateTranspose(self):
513        A = self.A
514        A.setOption(PETSc.Mat.Option.SYMMETRIC, True)
515        AT = PETSc.Mat().createTranspose(A)
516        x, y = A.createVecs()
517        xt, yt = AT.createVecs()
518        #
519        y.setRandom()
520        A.multTranspose(y, x)
521        y.copy(xt)
522        AT.mult(xt, yt)
523        self.assertTrue(yt.equal(x))
524        #
525        x.setRandom()
526        A.mult(x, y)
527        x.copy(yt)
528        AT.multTranspose(yt, xt)
529        self.assertTrue(xt.equal(y))
530        del A
531
532    def testConvert(self):
533        self.assertTrue(self.A.convert(PETSc.Mat.Type.AIJ,PETSc.Mat()).equal(self.A))
534        self.assertTrue(self.A.convert(PETSc.Mat.Type.BAIJ,PETSc.Mat()).equal(self.A))
535        self.assertTrue(self.A.convert(PETSc.Mat.Type.SBAIJ,PETSc.Mat()).equal(self.A))
536        self.assertTrue(self.A.convert(PETSc.Mat.Type.DENSE,PETSc.Mat()).equal(self.A))
537
538    def testShift(self):
539        old = self._getCtx().D.copy()
540        self.A.shift(-0.5)
541        D = self._getCtx().D
542        self.assertTrue(D.equal(old-0.5))
543
544    def testScale(self):
545        old = self._getCtx().D.copy()
546        self.A.scale(-0.5)
547        D = self._getCtx().D
548        self.assertTrue(D.equal(-0.5*old))
549
550
551# --------------------------------------------------------------------
552
553if __name__ == '__main__':
554    unittest.main()
555