xref: /petsc/src/binding/petsc4py/test/test_mat_py.py (revision 5f80ce2ab25dff0f4601e710601cbbcecf323266) 
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.transpose(product)
175            else:
176                raise RuntimeError('wrong configuration')
177            product.scale(self.s)
178        elif producttype == 'ABt':
179            if mat is A: # product = identity * B^T
180                B.transpose(product)
181            elif mat is B: # product = A * identity^T
182                A.copy(product, structure=True)
183            else:
184                raise RuntimeError('wrong configuration')
185            product.scale(self.s)
186        elif producttype == 'PtAP':
187            if mat is A: # product = P^T * identity * P
188                B.transposeMatMult(B, self.tmp)
189                self.tmp.copy(product, structure=True)
190                product.scale(self.s)
191            elif mat is B: # product = identity^T * A * identity
192                A.copy(product, structure=True)
193                product.scale(self.s**2)
194            else:
195                raise RuntimeError('wrong configuration')
196        elif producttype == 'RARt':
197            if mat is A: # product = R * identity * R^t
198                B.matTransposeMult(B, self.tmp)
199                self.tmp.copy(product, structure=True)
200                product.scale(self.s)
201            elif mat is B: # product = identity * A * identity^T
202                A.copy(product, structure=True)
203                product.scale(self.s**2)
204            else:
205                raise RuntimeError('wrong configuration')
206        elif producttype == 'ABC':
207            if mat is A: # product = identity * B * C
208                B.matMult(C, self.tmp)
209                self.tmp.copy(product, structure=True)
210            elif mat is B: # product = A * identity * C
211                A.matMult(C, self.tmp)
212                self.tmp.copy(product, structure=True)
213            elif mat is C: # product = A * B * identity
214                A.matMult(B, self.tmp)
215                self.tmp.copy(product, structure=True)
216            else:
217                raise RuntimeError('wrong configuration')
218            product.scale(self.s)
219        else:
220            raise RuntimeError('Product {} not implemented'.format(producttype))
221
222class Diagonal(Matrix):
223
224    def create(self, mat):
225        super(Diagonal,self).create(mat)
226        mat.setUp()
227        self.D = mat.createVecLeft()
228
229    def destroy(self, mat):
230        self.D.destroy()
231        super(Diagonal,self).destroy(mat)
232
233    def scale(self, mat, a):
234        self.D.scale(a)
235
236    def shift(self, mat, a):
237        self.D.shift(a)
238
239    def zeroEntries(self, mat):
240        self.D.zeroEntries()
241
242    def mult(self, mat, x, y):
243        y.pointwiseMult(x, self.D)
244
245    def duplicate(self, mat, op):
246        dmat = PETSc.Mat()
247        dctx = Diagonal()
248        dmat.createPython(mat.getSizes(), dctx, comm=mat.getComm())
249        dctx.D = self.D.duplicate()
250        if op == PETSc.Mat.DuplicateOption.COPY_VALUES:
251          self.D.copy(dctx.D)
252          dmat.setUp()
253        return dmat
254
255    def getDiagonal(self, mat, vd):
256        self.D.copy(vd)
257
258    def setDiagonal(self, mat, vd, im):
259        if isinstance (im, bool):
260            addv = im
261            if addv:
262                self.D.axpy(1, vd)
263            else:
264                vd.copy(self.D)
265        elif im == PETSc.InsertMode.INSERT_VALUES:
266            vd.copy(self.D)
267        elif im == PETSc.InsertMode.ADD_VALUES:
268            self.D.axpy(1, vd)
269        else:
270            raise ValueError('wrong InsertMode %d'% im)
271
272    def diagonalScale(self, mat, vl, vr):
273        if vl: self.D.pointwiseMult(self.D, vl)
274        if vr: self.D.pointwiseMult(self.D, vr)
275
276# --------------------------------------------------------------------
277
278class TestMatrix(unittest.TestCase):
279
280    COMM = PETSc.COMM_WORLD
281    PYMOD = __name__
282    PYCLS = 'Matrix'
283
284    def _getCtx(self):
285        return self.A.getPythonContext()
286
287    def setUp(self):
288        N = self.N = 13
289        self.A = PETSc.Mat()
290        if 0: # command line way
291            self.A.create(self.COMM)
292            self.A.setSizes([N,N])
293            self.A.setType('python')
294            OptDB = PETSc.Options(self.A)
295            OptDB['mat_python_type'] = '%s.%s' % (self.PYMOD,self.PYCLS)
296            self.A.setFromOptions()
297            self.A.setUp()
298            del OptDB['mat_python_type']
299            self.assertTrue(self._getCtx() is not None)
300        else: # python way
301            context = globals()[self.PYCLS]()
302            self.A.createPython([N,N], context, comm=self.COMM)
303            self.A.setUp()
304            self.assertTrue(self._getCtx() is context)
305            self.assertEqual(getrefcount(context), 3)
306            del context
307            self.assertEqual(getrefcount(self._getCtx()), 2)
308
309    def tearDown(self):
310        ctx = self.A.getPythonContext()
311        self.assertEqual(getrefcount(ctx), 3)
312        self.A.destroy() # XXX
313        self.A = None
314        self.assertEqual(getrefcount(ctx), 2)
315        #import gc,pprint; pprint.pprint(gc.get_referrers(ctx))
316
317    def testBasic(self):
318        ctx = self.A.getPythonContext()
319        self.assertTrue(self._getCtx() is ctx)
320        self.assertEqual(getrefcount(ctx), 3)
321
322    def testZeroEntries(self):
323        f = lambda : self.A.zeroEntries()
324        self.assertRaises(Exception, f)
325
326    def testMult(self):
327        x, y = self.A.createVecs()
328        f = lambda : self.A.mult(x, y)
329        self.assertRaises(Exception, f)
330
331    def testMultTranspose(self):
332        x, y = self.A.createVecs()
333        f = lambda : self.A.multTranspose(x, y)
334        self.assertRaises(Exception, f)
335
336    def testGetDiagonal(self):
337        d = self.A.createVecLeft()
338        f = lambda : self.A.getDiagonal(d)
339        self.assertRaises(Exception, f)
340
341    def testSetDiagonal(self):
342        d = self.A.createVecLeft()
343        f = lambda : self.A.setDiagonal(d)
344        self.assertRaises(Exception, f)
345
346    def testDiagonalScale(self):
347        x, y = self.A.createVecs()
348        f = lambda : self.A.diagonalScale(x, y)
349        self.assertRaises(Exception, f)
350
351    def testDuplicate(self):
352        f1 = lambda : self.A.duplicate(x, True)
353        f2 = lambda : self.A.duplicate(x, False)
354        self.assertRaises(Exception, f1)
355        self.assertRaises(Exception, f2)
356
357    def testSetVecType(self):
358        self.A.setVecType('mpi')
359        self.assertTrue('mpi' == self.A.getVecType())
360
361    def testH2Opus(self):
362        if not PETSc.Sys.hasExternalPackage("h2opus"):
363            return
364        if self.A.getComm().Get_size() > 1:
365            return
366        h = PETSc.Mat()
367
368        # need matrix vector and its transpose for norm estimation
369        AA = self.A.getPythonContext()
370        if not hasattr(AA,'mult'):
371            return
372        AA.multTranspose = AA.mult
373
374        # without coordinates
375        h.createH2OpusFromMat(self.A,leafsize=2)
376        h.assemble()
377        h.destroy()
378
379        # with coordinates
380        coords = numpy.linspace((1,2,3),(10,20,30),self.A.getSize()[0],dtype=PETSc.RealType)
381        h.createH2OpusFromMat(self.A,coords,leafsize=2)
382        h.assemble()
383
384        # test API
385        h.H2OpusOrthogonalize()
386        h.H2OpusCompress(1.e-1)
387
388        # Low-rank update
389        U = PETSc.Mat()
390        U.createDense([h.getSizes()[0],3],comm=h.getComm())
391        U.setUp()
392        U.setRandom()
393
394        he = PETSc.Mat()
395        h.convert('dense',he)
396        he.axpy(1.0, U.matTransposeMult(U))
397
398        h.H2OpusLowRankUpdate(U)
399        self.assertTrue(he.equal(h))
400
401
402        h.destroy()
403
404        del AA.multTranspose
405
406
407class TestScaledIdentity(TestMatrix):
408
409    PYCLS = 'ScaledIdentity'
410
411    def testMult(self):
412        s = self._getCtx().s
413        x, y = self.A.createVecs()
414        x.setRandom()
415        self.A.mult(x,y)
416        self.assertTrue(y.equal(s*x))
417
418    def testMultTransposeSymmKnown(self):
419        s = self._getCtx().s
420        x, y = self.A.createVecs()
421        x.setRandom()
422        self.A.setOption(PETSc.Mat.Option.SYMMETRIC, True)
423        self.A.multTranspose(x,y)
424        self.assertTrue(y.equal(s*x))
425        self.A.setOption(PETSc.Mat.Option.SYMMETRIC, False)
426        f = lambda : self.A.multTranspose(x, y)
427        self.assertRaises(Exception, f)
428
429    def testMultTransposeNewMeth(self):
430        s = self._getCtx().s
431        x, y = self.A.createVecs()
432        x.setRandom()
433        AA = self.A.getPythonContext()
434        AA.multTranspose = AA.mult
435        self.A.multTranspose(x,y)
436        del AA.multTranspose
437        self.assertTrue(y.equal(s*x))
438
439    def testGetDiagonal(self):
440        s = self._getCtx().s
441        d = self.A.createVecLeft()
442        o = d.duplicate()
443        o.set(s)
444        self.A.getDiagonal(d)
445        self.assertTrue(o.equal(d))
446
447    def testDuplicate(self):
448        B = self.A.duplicate(False)
449        self.assertTrue(B.getPythonContext().s == 2)
450        B = self.A.duplicate(True)
451        self.assertTrue(B.getPythonContext().s == self.A.getPythonContext().s)
452
453    def testMatMat(self):
454        s = self._getCtx().s
455        R = PETSc.Random().create(self.COMM)
456        R.setFromOptions()
457        A = PETSc.Mat().create(self.COMM)
458        A.setSizes(self.A.getSizes())
459        A.setType(PETSc.Mat.Type.AIJ)
460        A.setUp()
461        A.setRandom(R)
462        B = PETSc.Mat().create(self.COMM)
463        B.setSizes(self.A.getSizes())
464        B.setType(PETSc.Mat.Type.AIJ)
465        B.setUp()
466        B.setRandom(R)
467        I = PETSc.Mat().create(self.COMM)
468        I.setSizes(self.A.getSizes())
469        I.setType(PETSc.Mat.Type.AIJ)
470        I.setUp()
471        I.assemble()
472        I.shift(s)
473
474        self.assertTrue(self.A.matMult(A).equal(I.matMult(A)))
475        self.assertTrue(A.matMult(self.A).equal(A.matMult(I)))
476        if self.A.getComm().Get_size() == 1:
477            self.assertTrue(self.A.matTransposeMult(A).equal(I.matTransposeMult(A)))
478            self.assertTrue(A.matTransposeMult(self.A).equal(A.matTransposeMult(I)))
479        self.assertTrue(self.A.transposeMatMult(A).equal(I.transposeMatMult(A)))
480        self.assertTrue(A.transposeMatMult(self.A).equal(A.transposeMatMult(I)))
481        self.assertAlmostEqual((self.A.ptap(A) - I.ptap(A)).norm(), 0.0, places=5)
482        self.assertAlmostEqual((A.ptap(self.A) - A.ptap(I)).norm(), 0.0, places=5)
483        if self.A.getComm().Get_size() == 1:
484            self.assertAlmostEqual((self.A.rart(A) - I.rart(A)).norm(), 0.0, places=5)
485            self.assertAlmostEqual((A.rart(self.A) - A.rart(I)).norm(), 0.0, places=5)
486        self.assertAlmostEqual((self.A.matMatMult(A,B)-I.matMatMult(A,B)).norm(), 0.0, places=5)
487        self.assertAlmostEqual((A.matMatMult(self.A,B)-A.matMatMult(I,B)).norm(), 0.0, places=5)
488        self.assertAlmostEqual((A.matMatMult(B,self.A)-A.matMatMult(B,I)).norm(), 0.0, places=5)
489
490    def testShift(self):
491        sold = self._getCtx().s
492        self.A.shift(-0.5)
493        s = self._getCtx().s
494        self.assertTrue(s == sold - 0.5)
495
496    def testScale(self):
497        sold = self._getCtx().s
498        self.A.scale(-0.5)
499        s = self._getCtx().s
500        self.assertTrue(s == sold * -0.5)
501
502class TestDiagonal(TestMatrix):
503
504    PYCLS = 'Diagonal'
505
506    def setUp(self):
507        super(TestDiagonal, self).setUp()
508        D = self.A.createVecLeft()
509        s, e = D.getOwnershipRange()
510        for i in range(s, e):
511            D[i] = i+1
512        D.assemble()
513        self.A.setDiagonal(D)
514
515    def testZeroEntries(self):
516        self.A.zeroEntries()
517        D = self._getCtx().D
518        self.assertEqual(D.norm(), 0)
519
520    def testMult(self):
521        x, y = self.A.createVecs()
522        x.set(1)
523        self.A.mult(x,y)
524        self.assertTrue(y.equal(self._getCtx().D))
525
526    def testMultTransposeSymmKnown(self):
527        x, y = self.A.createVecs()
528        x.set(1)
529        self.A.setOption(PETSc.Mat.Option.SYMMETRIC, True)
530        self.A.multTranspose(x,y)
531        self.assertTrue(y.equal(self._getCtx().D))
532        self.A.setOption(PETSc.Mat.Option.SYMMETRIC, False)
533        f = lambda : self.A.multTranspose(x, y)
534        self.assertRaises(Exception, f)
535
536    def testMultTransposeNewMeth(self):
537        x, y = self.A.createVecs()
538        x.set(1)
539        AA = self.A.getPythonContext()
540        AA.multTranspose = AA.mult
541        self.A.multTranspose(x,y)
542        del AA.multTranspose
543        self.assertTrue(y.equal(self._getCtx().D))
544
545    def testDuplicate(self):
546        B = self.A.duplicate(False)
547        B = self.A.duplicate(True)
548        self.assertTrue(B.getPythonContext().D.equal(self.A.getPythonContext().D))
549
550    def testGetDiagonal(self):
551        d = self.A.createVecLeft()
552        self.A.getDiagonal(d)
553        self.assertTrue(d.equal(self._getCtx().D))
554
555    def testSetDiagonal(self):
556        d = self.A.createVecLeft()
557        d.setRandom()
558        self.A.setDiagonal(d)
559        self.assertTrue(d.equal(self._getCtx().D))
560
561    def testDiagonalScale(self):
562        x, y = self.A.createVecs()
563        x.set(2)
564        y.set(3)
565        old = self._getCtx().D.copy()
566        self.A.diagonalScale(x, y)
567        D = self._getCtx().D
568        self.assertTrue(D.equal(old*6))
569
570    def testCreateTranspose(self):
571        A = self.A
572        A.setOption(PETSc.Mat.Option.SYMMETRIC, True)
573        AT = PETSc.Mat().createTranspose(A)
574        x, y = A.createVecs()
575        xt, yt = AT.createVecs()
576        #
577        y.setRandom()
578        A.multTranspose(y, x)
579        y.copy(xt)
580        AT.mult(xt, yt)
581        self.assertTrue(yt.equal(x))
582        #
583        x.setRandom()
584        A.mult(x, y)
585        x.copy(yt)
586        AT.multTranspose(yt, xt)
587        self.assertTrue(xt.equal(y))
588        del A
589
590    def testConvert(self):
591        self.assertTrue(self.A.convert(PETSc.Mat.Type.AIJ,PETSc.Mat()).equal(self.A))
592        self.assertTrue(self.A.convert(PETSc.Mat.Type.BAIJ,PETSc.Mat()).equal(self.A))
593        self.assertTrue(self.A.convert(PETSc.Mat.Type.SBAIJ,PETSc.Mat()).equal(self.A))
594        self.assertTrue(self.A.convert(PETSc.Mat.Type.DENSE,PETSc.Mat()).equal(self.A))
595
596    def testShift(self):
597        old = self._getCtx().D.copy()
598        self.A.shift(-0.5)
599        D = self._getCtx().D
600        self.assertTrue(D.equal(old-0.5))
601
602    def testScale(self):
603        old = self._getCtx().D.copy()
604        self.A.scale(-0.5)
605        D = self._getCtx().D
606        self.assertTrue(D.equal(-0.5*old))
607
608
609# --------------------------------------------------------------------
610
611if __name__ == '__main__':
612    unittest.main()
613