xref: /petsc/src/binding/petsc4py/test/test_mat_py.py (revision a69119a591a03a9d906b29c0a4e9802e4d7c9795) 
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        self.assertEqual(getrefcount(ctx), 2)
317        #import gc,pprint; pprint.pprint(gc.get_referrers(ctx))
318
319    def testBasic(self):
320        ctx = self.A.getPythonContext()
321        self.assertTrue(self._getCtx() is ctx)
322        self.assertEqual(getrefcount(ctx), 3)
323
324    def testZeroEntries(self):
325        f = lambda : self.A.zeroEntries()
326        self.assertRaises(Exception, f)
327
328    def testMult(self):
329        x, y = self.A.createVecs()
330        f = lambda : self.A.mult(x, y)
331        self.assertRaises(Exception, f)
332
333    def testMultTranspose(self):
334        x, y = self.A.createVecs()
335        f = lambda : self.A.multTranspose(x, y)
336        self.assertRaises(Exception, f)
337
338    def testGetDiagonal(self):
339        d = self.A.createVecLeft()
340        f = lambda : self.A.getDiagonal(d)
341        self.assertRaises(Exception, f)
342
343    def testSetDiagonal(self):
344        d = self.A.createVecLeft()
345        f = lambda : self.A.setDiagonal(d)
346        self.assertRaises(Exception, f)
347
348    def testDiagonalScale(self):
349        x, y = self.A.createVecs()
350        f = lambda : self.A.diagonalScale(x, y)
351        self.assertRaises(Exception, f)
352
353    def testDuplicate(self):
354        f1 = lambda : self.A.duplicate(x, True)
355        f2 = lambda : self.A.duplicate(x, False)
356        self.assertRaises(Exception, f1)
357        self.assertRaises(Exception, f2)
358
359    def testSetVecType(self):
360        self.A.setVecType('mpi')
361        self.assertTrue('mpi' == self.A.getVecType())
362
363    def testH2Opus(self):
364        if not PETSc.Sys.hasExternalPackage("h2opus"):
365            return
366        if self.A.getComm().Get_size() > 1:
367            return
368        h = PETSc.Mat()
369
370        # need matrix vector and its transpose for norm estimation
371        AA = self.A.getPythonContext()
372        if not hasattr(AA,'mult'):
373            return
374        AA.multTranspose = AA.mult
375
376        # without coordinates
377        h.createH2OpusFromMat(self.A,leafsize=2)
378        h.assemble()
379        h.destroy()
380
381        # with coordinates
382        coords = numpy.linspace((1,2,3),(10,20,30),self.A.getSize()[0],dtype=PETSc.RealType)
383        h.createH2OpusFromMat(self.A,coords,leafsize=2)
384        h.assemble()
385
386        # test API
387        h.H2OpusOrthogonalize()
388        h.H2OpusCompress(1.e-1)
389
390        # Low-rank update
391        U = PETSc.Mat()
392        U.createDense([h.getSizes()[0],3],comm=h.getComm())
393        U.setUp()
394        U.setRandom()
395
396        he = PETSc.Mat()
397        h.convert('dense',he)
398        he.axpy(1.0, U.matTransposeMult(U))
399
400        h.H2OpusLowRankUpdate(U)
401        self.assertTrue(he.equal(h))
402
403
404        h.destroy()
405
406        del AA.multTranspose
407
408    def testGetType(self):
409        ctx = self.A.getPythonContext()
410        pytype = "{0}.{1}".format(ctx.__module__, type(ctx).__name__)
411        self.assertTrue(self.A.getPythonType() == pytype)
412
413class TestScaledIdentity(TestMatrix):
414
415    PYCLS = 'ScaledIdentity'
416
417    def testMult(self):
418        s = self._getCtx().s
419        x, y = self.A.createVecs()
420        x.setRandom()
421        self.A.mult(x,y)
422        self.assertTrue(y.equal(s*x))
423
424    def testMultTransposeSymmKnown(self):
425        s = self._getCtx().s
426        x, y = self.A.createVecs()
427        x.setRandom()
428        self.A.setOption(PETSc.Mat.Option.SYMMETRIC, True)
429        self.A.multTranspose(x,y)
430        self.assertTrue(y.equal(s*x))
431        self.A.setOption(PETSc.Mat.Option.SYMMETRIC, False)
432        f = lambda : self.A.multTranspose(x, y)
433        self.assertRaises(Exception, f)
434
435    def testMultTransposeNewMeth(self):
436        s = self._getCtx().s
437        x, y = self.A.createVecs()
438        x.setRandom()
439        AA = self.A.getPythonContext()
440        AA.multTranspose = AA.mult
441        self.A.multTranspose(x,y)
442        del AA.multTranspose
443        self.assertTrue(y.equal(s*x))
444
445    def testGetDiagonal(self):
446        s = self._getCtx().s
447        d = self.A.createVecLeft()
448        o = d.duplicate()
449        o.set(s)
450        self.A.getDiagonal(d)
451        self.assertTrue(o.equal(d))
452
453    def testDuplicate(self):
454        B = self.A.duplicate(False)
455        self.assertTrue(B.getPythonContext().s == 2)
456        B = self.A.duplicate(True)
457        self.assertTrue(B.getPythonContext().s == self.A.getPythonContext().s)
458
459    def testMatMat(self):
460        s = self._getCtx().s
461        R = PETSc.Random().create(self.COMM)
462        R.setFromOptions()
463        A = PETSc.Mat().create(self.COMM)
464        A.setSizes(self.A.getSizes())
465        A.setType(PETSc.Mat.Type.AIJ)
466        A.setUp()
467        A.setRandom(R)
468        B = PETSc.Mat().create(self.COMM)
469        B.setSizes(self.A.getSizes())
470        B.setType(PETSc.Mat.Type.AIJ)
471        B.setUp()
472        B.setRandom(R)
473        I = PETSc.Mat().create(self.COMM)
474        I.setSizes(self.A.getSizes())
475        I.setType(PETSc.Mat.Type.AIJ)
476        I.setUp()
477        I.assemble()
478        I.shift(s)
479
480        self.assertTrue(self.A.matMult(A).equal(I.matMult(A)))
481        self.assertTrue(A.matMult(self.A).equal(A.matMult(I)))
482        if self.A.getComm().Get_size() == 1:
483            self.assertTrue(self.A.matTransposeMult(A).equal(I.matTransposeMult(A)))
484            self.assertTrue(A.matTransposeMult(self.A).equal(A.matTransposeMult(I)))
485        self.assertTrue(self.A.transposeMatMult(A).equal(I.transposeMatMult(A)))
486        self.assertTrue(A.transposeMatMult(self.A).equal(A.transposeMatMult(I)))
487        self.assertAlmostEqual((self.A.ptap(A) - I.ptap(A)).norm(), 0.0, places=5)
488        self.assertAlmostEqual((A.ptap(self.A) - A.ptap(I)).norm(), 0.0, places=5)
489        if self.A.getComm().Get_size() == 1:
490            self.assertAlmostEqual((self.A.rart(A) - I.rart(A)).norm(), 0.0, places=5)
491            self.assertAlmostEqual((A.rart(self.A) - A.rart(I)).norm(), 0.0, places=5)
492        self.assertAlmostEqual((self.A.matMatMult(A,B)-I.matMatMult(A,B)).norm(), 0.0, places=5)
493        self.assertAlmostEqual((A.matMatMult(self.A,B)-A.matMatMult(I,B)).norm(), 0.0, places=5)
494        self.assertAlmostEqual((A.matMatMult(B,self.A)-A.matMatMult(B,I)).norm(), 0.0, places=5)
495
496    def testShift(self):
497        sold = self._getCtx().s
498        self.A.shift(-0.5)
499        s = self._getCtx().s
500        self.assertTrue(s == sold - 0.5)
501
502    def testScale(self):
503        sold = self._getCtx().s
504        self.A.scale(-0.5)
505        s = self._getCtx().s
506        self.assertTrue(s == sold * -0.5)
507
508class TestDiagonal(TestMatrix):
509
510    PYCLS = 'Diagonal'
511
512    def setUp(self):
513        super(TestDiagonal, self).setUp()
514        D = self.A.createVecLeft()
515        s, e = D.getOwnershipRange()
516        for i in range(s, e):
517            D[i] = i+1
518        D.assemble()
519        self.A.setDiagonal(D)
520
521    def testZeroEntries(self):
522        self.A.zeroEntries()
523        D = self._getCtx().D
524        self.assertEqual(D.norm(), 0)
525
526    def testMult(self):
527        x, y = self.A.createVecs()
528        x.set(1)
529        self.A.mult(x,y)
530        self.assertTrue(y.equal(self._getCtx().D))
531
532    def testMultTransposeSymmKnown(self):
533        x, y = self.A.createVecs()
534        x.set(1)
535        self.A.setOption(PETSc.Mat.Option.SYMMETRIC, True)
536        self.A.multTranspose(x,y)
537        self.assertTrue(y.equal(self._getCtx().D))
538        self.A.setOption(PETSc.Mat.Option.SYMMETRIC, False)
539        f = lambda : self.A.multTranspose(x, y)
540        self.assertRaises(Exception, f)
541
542    def testMultTransposeNewMeth(self):
543        x, y = self.A.createVecs()
544        x.set(1)
545        AA = self.A.getPythonContext()
546        AA.multTranspose = AA.mult
547        self.A.multTranspose(x,y)
548        del AA.multTranspose
549        self.assertTrue(y.equal(self._getCtx().D))
550
551    def testDuplicate(self):
552        B = self.A.duplicate(False)
553        B = self.A.duplicate(True)
554        self.assertTrue(B.getPythonContext().D.equal(self.A.getPythonContext().D))
555
556    def testGetDiagonal(self):
557        d = self.A.createVecLeft()
558        self.A.getDiagonal(d)
559        self.assertTrue(d.equal(self._getCtx().D))
560
561    def testSetDiagonal(self):
562        d = self.A.createVecLeft()
563        d.setRandom()
564        self.A.setDiagonal(d)
565        self.assertTrue(d.equal(self._getCtx().D))
566
567    def testDiagonalScale(self):
568        x, y = self.A.createVecs()
569        x.set(2)
570        y.set(3)
571        old = self._getCtx().D.copy()
572        self.A.diagonalScale(x, y)
573        D = self._getCtx().D
574        self.assertTrue(D.equal(old*6))
575
576    def testCreateTranspose(self):
577        A = self.A
578        A.setOption(PETSc.Mat.Option.SYMMETRIC, True)
579        AT = PETSc.Mat().createTranspose(A)
580        x, y = A.createVecs()
581        xt, yt = AT.createVecs()
582        #
583        y.setRandom()
584        A.multTranspose(y, x)
585        y.copy(xt)
586        AT.mult(xt, yt)
587        self.assertTrue(yt.equal(x))
588        #
589        x.setRandom()
590        A.mult(x, y)
591        x.copy(yt)
592        AT.multTranspose(yt, xt)
593        self.assertTrue(xt.equal(y))
594        del A
595
596    def testConvert(self):
597        self.assertTrue(self.A.convert(PETSc.Mat.Type.AIJ,PETSc.Mat()).equal(self.A))
598        self.assertTrue(self.A.convert(PETSc.Mat.Type.BAIJ,PETSc.Mat()).equal(self.A))
599        self.assertTrue(self.A.convert(PETSc.Mat.Type.SBAIJ,PETSc.Mat()).equal(self.A))
600        self.assertTrue(self.A.convert(PETSc.Mat.Type.DENSE,PETSc.Mat()).equal(self.A))
601
602    def testShift(self):
603        old = self._getCtx().D.copy()
604        self.A.shift(-0.5)
605        D = self._getCtx().D
606        self.assertTrue(D.equal(old-0.5))
607
608    def testScale(self):
609        old = self._getCtx().D.copy()
610        self.A.scale(-0.5)
611        D = self._getCtx().D
612        self.assertTrue(D.equal(-0.5*old))
613
614
615# --------------------------------------------------------------------
616
617if __name__ == '__main__':
618    unittest.main()
619