xref: /petsc/src/binding/petsc4py/test/test_mat_py.py (revision 22fceea1769ca91bdb9988b063eaa3e47b647107) 
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
332class TestScaledIdentity(TestMatrix):
333
334    PYCLS = 'ScaledIdentity'
335
336    def testMult(self):
337        s = self._getCtx().s
338        x, y = self.A.createVecs()
339        x.setRandom()
340        self.A.mult(x,y)
341        self.assertTrue(y.equal(s*x))
342
343    def testMultTransposeSymmKnown(self):
344        s = self._getCtx().s
345        x, y = self.A.createVecs()
346        x.setRandom()
347        self.A.setOption(PETSc.Mat.Option.SYMMETRIC, True)
348        self.A.multTranspose(x,y)
349        self.assertTrue(y.equal(s*x))
350        self.A.setOption(PETSc.Mat.Option.SYMMETRIC, False)
351        f = lambda : self.A.multTranspose(x, y)
352        self.assertRaises(Exception, f)
353
354    def testMultTransposeNewMeth(self):
355        s = self._getCtx().s
356        x, y = self.A.createVecs()
357        x.setRandom()
358        AA = self.A.getPythonContext()
359        AA.multTranspose = AA.mult
360        self.A.multTranspose(x,y)
361        del AA.multTranspose
362        self.assertTrue(y.equal(s*x))
363
364    def testGetDiagonal(self):
365        s = self._getCtx().s
366        d = self.A.createVecLeft()
367        o = d.duplicate()
368        o.set(s)
369        self.A.getDiagonal(d)
370        self.assertTrue(o.equal(d))
371
372    def testMatMat(self):
373        s = self._getCtx().s
374        R = PETSc.Random().create(self.COMM)
375        R.setFromOptions()
376        A = PETSc.Mat().create(self.COMM)
377        A.setSizes(self.A.getSizes())
378        A.setType(PETSc.Mat.Type.AIJ)
379        A.setUp()
380        A.setRandom(R)
381        B = PETSc.Mat().create(self.COMM)
382        B.setSizes(self.A.getSizes())
383        B.setType(PETSc.Mat.Type.AIJ)
384        B.setUp()
385        B.setRandom(R)
386        I = PETSc.Mat().create(self.COMM)
387        I.setSizes(self.A.getSizes())
388        I.setType(PETSc.Mat.Type.AIJ)
389        I.setUp()
390        I.assemble()
391        I.shift(s)
392
393        self.assertTrue(self.A.matMult(A).equal(I.matMult(A)))
394        self.assertTrue(A.matMult(self.A).equal(A.matMult(I)))
395        if self.A.getComm().Get_size() == 1:
396            self.assertTrue(self.A.matTransposeMult(A).equal(I.matTransposeMult(A)))
397            self.assertTrue(A.matTransposeMult(self.A).equal(A.matTransposeMult(I)))
398        self.assertTrue(self.A.transposeMatMult(A).equal(I.transposeMatMult(A)))
399        self.assertTrue(A.transposeMatMult(self.A).equal(A.transposeMatMult(I)))
400        self.assertAlmostEqual((self.A.ptap(A) - I.ptap(A)).norm(), 0.0, places=5)
401        self.assertAlmostEqual((A.ptap(self.A) - A.ptap(I)).norm(), 0.0, places=5)
402        if self.A.getComm().Get_size() == 1:
403            self.assertAlmostEqual((self.A.rart(A) - I.rart(A)).norm(), 0.0, places=5)
404            self.assertAlmostEqual((A.rart(self.A) - A.rart(I)).norm(), 0.0, places=5)
405        self.assertAlmostEqual((self.A.matMatMult(A,B)-I.matMatMult(A,B)).norm(), 0.0, places=5)
406        self.assertAlmostEqual((A.matMatMult(self.A,B)-A.matMatMult(I,B)).norm(), 0.0, places=5)
407        self.assertAlmostEqual((A.matMatMult(B,self.A)-A.matMatMult(B,I)).norm(), 0.0, places=5)
408
409    def testH2Opus(self):
410        if not PETSc.Sys.hasExternalPackage("h2opus"):
411            return
412        if self.A.getComm().Get_size() > 1:
413            return
414        h = PETSc.Mat()
415
416        # need transpose operation for norm estimation
417        AA = self.A.getPythonContext()
418        AA.multTranspose = AA.mult
419
420        # without coordinates
421        h.createH2OpusFromMat(self.A,leafsize=2)
422        h.assemble()
423        h.destroy()
424
425        # with coordinates
426        coords = numpy.linspace((1,2,3),(10,20,30),self.A.getSize()[0],dtype=PETSc.RealType)
427        h.createH2OpusFromMat(self.A,coords,leafsize=2)
428        h.assemble()
429        h.destroy()
430
431        del AA.multTranspose
432
433    def testShift(self):
434        sold = self._getCtx().s
435        self.A.shift(-0.5)
436        s = self._getCtx().s
437        self.assertTrue(s == sold - 0.5)
438
439    def testScale(self):
440        sold = self._getCtx().s
441        self.A.scale(-0.5)
442        s = self._getCtx().s
443        self.assertTrue(s == sold * -0.5)
444
445class TestDiagonal(TestMatrix):
446
447    PYCLS = 'Diagonal'
448
449    def setUp(self):
450        super(TestDiagonal, self).setUp()
451        D = self.A.createVecLeft()
452        s, e = D.getOwnershipRange()
453        for i in range(s, e):
454            D[i] = i+1
455        D.assemble()
456        self.A.setDiagonal(D)
457
458    def testZeroEntries(self):
459        self.A.zeroEntries()
460        D = self._getCtx().D
461        self.assertEqual(D.norm(), 0)
462
463    def testMult(self):
464        x, y = self.A.createVecs()
465        x.set(1)
466        self.A.mult(x,y)
467        self.assertTrue(y.equal(self._getCtx().D))
468
469    def testMultTransposeSymmKnown(self):
470        x, y = self.A.createVecs()
471        x.set(1)
472        self.A.setOption(PETSc.Mat.Option.SYMMETRIC, True)
473        self.A.multTranspose(x,y)
474        self.assertTrue(y.equal(self._getCtx().D))
475        self.A.setOption(PETSc.Mat.Option.SYMMETRIC, False)
476        f = lambda : self.A.multTranspose(x, y)
477        self.assertRaises(Exception, f)
478
479    def testMultTransposeNewMeth(self):
480        x, y = self.A.createVecs()
481        x.set(1)
482        AA = self.A.getPythonContext()
483        AA.multTranspose = AA.mult
484        self.A.multTranspose(x,y)
485        del AA.multTranspose
486        self.assertTrue(y.equal(self._getCtx().D))
487
488    def testGetDiagonal(self):
489        d = self.A.createVecLeft()
490        self.A.getDiagonal(d)
491        self.assertTrue(d.equal(self._getCtx().D))
492
493    def testSetDiagonal(self):
494        d = self.A.createVecLeft()
495        d.setRandom()
496        self.A.setDiagonal(d)
497        self.assertTrue(d.equal(self._getCtx().D))
498
499    def testDiagonalScale(self):
500        x, y = self.A.createVecs()
501        x.set(2)
502        y.set(3)
503        old = self._getCtx().D.copy()
504        self.A.diagonalScale(x, y)
505        D = self._getCtx().D
506        self.assertTrue(D.equal(old*6))
507
508    def testCreateTranspose(self):
509        A = self.A
510        A.setOption(PETSc.Mat.Option.SYMMETRIC, True)
511        AT = PETSc.Mat().createTranspose(A)
512        x, y = A.createVecs()
513        xt, yt = AT.createVecs()
514        #
515        y.setRandom()
516        A.multTranspose(y, x)
517        y.copy(xt)
518        AT.mult(xt, yt)
519        self.assertTrue(yt.equal(x))
520        #
521        x.setRandom()
522        A.mult(x, y)
523        x.copy(yt)
524        AT.multTranspose(yt, xt)
525        self.assertTrue(xt.equal(y))
526        del A
527
528    def testConvert(self):
529        self.assertTrue(self.A.convert(PETSc.Mat.Type.AIJ,PETSc.Mat()).equal(self.A))
530        self.assertTrue(self.A.convert(PETSc.Mat.Type.BAIJ,PETSc.Mat()).equal(self.A))
531        self.assertTrue(self.A.convert(PETSc.Mat.Type.SBAIJ,PETSc.Mat()).equal(self.A))
532        self.assertTrue(self.A.convert(PETSc.Mat.Type.DENSE,PETSc.Mat()).equal(self.A))
533
534    def testShift(self):
535        old = self._getCtx().D.copy()
536        self.A.shift(-0.5)
537        D = self._getCtx().D
538        self.assertTrue(D.equal(old-0.5))
539
540    def testScale(self):
541        old = self._getCtx().D.copy()
542        self.A.scale(-0.5)
543        D = self._getCtx().D
544        self.assertTrue(D.equal(-0.5*old))
545
546
547# --------------------------------------------------------------------
548
549if __name__ == '__main__':
550    unittest.main()
551