xref: /petsc/src/binding/petsc4py/test/test_mat_py.py (revision 0aed07b9805c9616851d69fcfe7cc3cc71b23e26) 
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 Identity(Matrix):
18
19    def mult(self, mat, x, y):
20        x.copy(y)
21
22    def getDiagonal(self, mat, vd):
23        vd.set(1)
24
25class Diagonal(Matrix):
26
27    def create(self, mat):
28        super(Diagonal,self).create(mat)
29        mat.setUp()
30        self.D = mat.createVecLeft()
31
32    def destroy(self, mat):
33        self.D.destroy()
34        super(Diagonal,self).destroy(mat)
35
36    def scale(self, mat, a):
37        self.D.scale(a)
38
39    def shift(self, mat, a):
40        self.D.shift(a)
41
42    def zeroEntries(self, mat):
43        self.D.zeroEntries()
44
45    def mult(self, mat, x, y):
46        y.pointwiseMult(x, self.D)
47
48    def getDiagonal(self, mat, vd):
49        self.D.copy(vd)
50
51    def setDiagonal(self, mat, vd, im):
52        if isinstance (im, bool):
53            addv = im
54            if addv:
55                self.D.axpy(1, vd)
56            else:
57                vd.copy(self.D)
58        elif im == PETSc.InsertMode.INSERT_VALUES:
59            vd.copy(self.D)
60        elif im == PETSc.InsertMode.ADD_VALUES:
61            self.D.axpy(1, vd)
62        else:
63            raise ValueError('wrong InsertMode %d'% im)
64
65    def diagonalScale(self, mat, vl, vr):
66        if vl: self.D.pointwiseMult(self.D, vl)
67        if vr: self.D.pointwiseMult(self.D, vr)
68
69# --------------------------------------------------------------------
70
71class TestMatrix(unittest.TestCase):
72
73    COMM = PETSc.COMM_WORLD
74    PYMOD = __name__
75    PYCLS = 'Matrix'
76
77    def _getCtx(self):
78        return self.A.getPythonContext()
79
80    def setUp(self):
81        N = self.N = 10
82        self.A = PETSc.Mat()
83        if 0: # command line way
84            self.A.create(self.COMM)
85            self.A.setSizes([N,N])
86            self.A.setType('python')
87            OptDB = PETSc.Options(self.A)
88            OptDB['mat_python_type'] = '%s.%s' % (self.PYMOD,self.PYCLS)
89            self.A.setFromOptions()
90            self.A.setUp()
91            del OptDB['mat_python_type']
92            self.assertTrue(self._getCtx() is not None)
93        else: # python way
94            context = globals()[self.PYCLS]()
95            self.A.createPython([N,N], context, comm=self.COMM)
96            self.A.setUp()
97            self.assertTrue(self._getCtx() is context)
98            self.assertEqual(getrefcount(context), 3)
99            del context
100            self.assertEqual(getrefcount(self._getCtx()), 2)
101
102    def tearDown(self):
103        ctx = self.A.getPythonContext()
104        self.assertEqual(getrefcount(ctx), 3)
105        self.A.destroy() # XXX
106        self.A = None
107        self.assertEqual(getrefcount(ctx), 2)
108        #import gc,pprint; pprint.pprint(gc.get_referrers(ctx))
109
110    def testBasic(self):
111        ctx = self.A.getPythonContext()
112        self.assertTrue(self._getCtx() is ctx)
113        self.assertEqual(getrefcount(ctx), 3)
114
115    def testZeroEntries(self):
116        f = lambda : self.A.zeroEntries()
117        self.assertRaises(Exception, f)
118
119    def testMult(self):
120        x, y = self.A.createVecs()
121        f = lambda : self.A.mult(x, y)
122        self.assertRaises(Exception, f)
123
124    def testMultTranspose(self):
125        x, y = self.A.createVecs()
126        f = lambda : self.A.multTranspose(x, y)
127        self.assertRaises(Exception, f)
128
129    def testGetDiagonal(self):
130        d = self.A.createVecLeft()
131        f = lambda : self.A.getDiagonal(d)
132        self.assertRaises(Exception, f)
133
134    def testSetDiagonal(self):
135        d = self.A.createVecLeft()
136        f = lambda : self.A.setDiagonal(d)
137        self.assertRaises(Exception, f)
138
139    def testDiagonalScale(self):
140        x, y = self.A.createVecs()
141        f = lambda : self.A.diagonalScale(x, y)
142        self.assertRaises(Exception, f)
143
144
145class TestIdentity(TestMatrix):
146
147    PYCLS = 'Identity'
148
149    def testMult(self):
150        x, y = self.A.createVecs()
151        x.setRandom()
152        self.A.mult(x,y)
153        self.assertTrue(y.equal(x))
154
155    def testMultTransposeSymmKnown(self):
156        x, y = self.A.createVecs()
157        x.setRandom()
158        self.A.setOption(PETSc.Mat.Option.SYMMETRIC, True)
159        self.A.multTranspose(x,y)
160        self.assertTrue(y.equal(x))
161        self.A.setOption(PETSc.Mat.Option.SYMMETRIC, False)
162        f = lambda : self.A.multTranspose(x, y)
163        self.assertRaises(Exception, f)
164
165    def testMultTransposeNewMeth(self):
166        x, y = self.A.createVecs()
167        x.setRandom()
168        AA = self.A.getPythonContext()
169        AA.multTranspose = AA.mult
170        self.A.multTranspose(x,y)
171        del AA.multTranspose
172        self.assertTrue(y.equal(x))
173
174    def testGetDiagonal(self):
175        d = self.A.createVecLeft()
176        o = d.duplicate()
177        o.set(1)
178        self.A.getDiagonal(d)
179        self.assertTrue(o.equal(d))
180
181    def testH2Opus(self):
182        if not PETSc.Sys.hasExternalPackage("h2opus"):
183            return
184        h = PETSc.Mat()
185
186        # need transpose operation for norm estimation
187        AA = self.A.getPythonContext()
188        AA.multTranspose = AA.mult
189
190        # without coordinates
191        h.createH2OpusFromMat(self.A,leafsize=2)
192        h.assemble()
193        h.destroy()
194
195        # with coordinates
196        coords = numpy.linspace((1,2,3),(10,20,30),self.A.getSize()[0],dtype=PETSc.RealType)
197        h.createH2OpusFromMat(self.A,coords,leafsize=2)
198        h.assemble()
199        h.destroy()
200
201        del AA.multTranspose
202
203class TestDiagonal(TestMatrix):
204
205    PYCLS = 'Diagonal'
206
207    def setUp(self):
208        super(TestDiagonal, self).setUp()
209        D = self.A.createVecLeft()
210        s, e = D.getOwnershipRange()
211        for i in range(s, e):
212            D[i] = i+1
213        D.assemble()
214        self.A.setDiagonal(D)
215
216
217    def testZeroEntries(self):
218        self.A.zeroEntries()
219        D = self._getCtx().D
220        self.assertEqual(D.norm(), 0)
221
222    def testMult(self):
223        x, y = self.A.createVecs()
224        x.set(1)
225        self.A.mult(x,y)
226        self.assertTrue(y.equal(self._getCtx().D))
227
228    def testMultTransposeSymmKnown(self):
229        x, y = self.A.createVecs()
230        x.set(1)
231        self.A.setOption(PETSc.Mat.Option.SYMMETRIC, True)
232        self.A.multTranspose(x,y)
233        self.assertTrue(y.equal(self._getCtx().D))
234        self.A.setOption(PETSc.Mat.Option.SYMMETRIC, False)
235        f = lambda : self.A.multTranspose(x, y)
236        self.assertRaises(Exception, f)
237
238    def testMultTransposeNewMeth(self):
239        x, y = self.A.createVecs()
240        x.set(1)
241        AA = self.A.getPythonContext()
242        AA.multTranspose = AA.mult
243        self.A.multTranspose(x,y)
244        del AA.multTranspose
245        self.assertTrue(y.equal(self._getCtx().D))
246
247    def testGetDiagonal(self):
248        d = self.A.createVecLeft()
249        self.A.getDiagonal(d)
250        self.assertTrue(d.equal(self._getCtx().D))
251
252    def testSetDiagonal(self):
253        d = self.A.createVecLeft()
254        d.setRandom()
255        self.A.setDiagonal(d)
256        self.assertTrue(d.equal(self._getCtx().D))
257
258    def testDiagonalScale(self):
259        x, y = self.A.createVecs()
260        x.set(2)
261        y.set(3)
262        old = self._getCtx().D.copy()
263        self.A.diagonalScale(x, y)
264        D = self._getCtx().D
265        self.assertTrue(D.equal(old*6))
266
267    def testCreateTranspose(self):
268        A = self.A
269        A.setOption(PETSc.Mat.Option.SYMMETRIC, True)
270        AT = PETSc.Mat().createTranspose(A)
271        x, y = A.createVecs()
272        xt, yt = AT.createVecs()
273        #
274        y.setRandom()
275        A.multTranspose(y, x)
276        y.copy(xt)
277        AT.mult(xt, yt)
278        self.assertTrue(yt.equal(x))
279        #
280        x.setRandom()
281        A.mult(x, y)
282        x.copy(yt)
283        AT.multTranspose(yt, xt)
284        self.assertTrue(xt.equal(y))
285        del A
286
287
288# --------------------------------------------------------------------
289
290if __name__ == '__main__':
291    unittest.main()
292