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