xref: /petsc/src/binding/petsc4py/test/test_mat_dense.py (revision f73bfc444f01cd53364f186d523b473a6ee5b07a) ! 
1from petsc4py import PETSc
2import unittest
3
4import numpy as np
5
6
7def mkdata(comm, m, N, bs):
8    start = m * comm.rank
9    end = start + m
10    idt = PETSc.IntType
11    sdt = PETSc.ScalarType
12    rows = np.array(range(start, end), dtype=idt)
13    cols = np.array(range(N), dtype=idt)
14    vals = np.array(range(m * N * bs * bs), dtype=sdt)
15    vals.shape = (-1, bs, bs)
16    return rows, cols, vals
17
18
19class BaseTestMatAnyDense:
20    COMM = PETSc.COMM_NULL
21    GRID = 0, 0
22    BSIZE = None
23    TYPE = PETSc.Mat.Type.DENSE
24
25    def setUp(self):
26        COMM = self.COMM
27        GM, GN = self.GRID
28        BS = self.BSIZE  # or 1
29        #
30        self.A = PETSc.Mat().create(comm=COMM)
31        bs = BS or 1
32        m, N = GM, GN
33        rowsz = (m * bs, None)
34        colsz = (None, N * bs)
35        self.A.setSizes([rowsz, colsz], BS)
36        self.A.setType(self.TYPE)
37
38    def tearDown(self):
39        self.A.destroy()
40        self.A = None
41        PETSc.garbage_cleanup()
42
43    def testSetValues(self):
44        self._preallocate()
45        r, c, v = self._set_values()
46        self.A.assemble()
47        self._chk_array(self.A, r, c, v)
48        r, c, v = self._set_values()
49        self.A.assemble()
50        self._chk_array(self.A, r, c, v)
51
52    def testGetDiagonalBlock(self):
53        M, N = self.A.getSize()
54        # only for square matrices
55        if M != N:
56            return
57        self._preallocate()
58        self._set_values()
59        self.A.assemble()
60        B = self.A.getDiagonalBlock()
61        self.assertEqual(self.A.getLocalSize(), B.getSize())
62        B.destroy()
63
64    def testSubMatrix(self):
65        self._preallocate()
66        self._set_values()
67        self.A.assemble()
68        B = self.A.getDenseSubMatrix()
69        X = B.copy()
70        self.A.restoreDenseSubMatrix(B)
71        self.assertTrue(self.A.equal(X))
72        X.destroy()
73        M, N = self.A.getSize()
74        rst = min(1,M)
75        cst = min(2,N)
76        ren = min(3,M)
77        cen = min(5,N)
78        B = self.A.getDenseSubMatrix(rst, ren, cst, cen)
79        self.assertTrue(B.getSize(), (ren - rst, cen - cst))
80        self.A.restoreDenseSubMatrix(B)
81        self.assertFalse(B)
82        B = self.A.getDenseSubMatrix(cbegin=cst, rbegin=rst, cend=cen, rend=ren)
83        self.assertTrue(B.getSize(), (ren - rst, cen - cst))
84        self.A.restoreDenseSubMatrix(B)
85        self.assertFalse(B)
86
87    def testColumnVec(self):
88        self._preallocate()
89        self._set_values()
90        self.A.assemble()
91        x = self.A.createVecLeft()
92        x.setRandom()
93        for i in range(self.A.getSize()[1]):
94            c = self.A.getDenseColumnVec(i)
95            x.copy(c)
96            self.A.restoreDenseColumnVec(i, V=c)
97            self.assertFalse(c)
98            c = self.A.getDenseColumnVec(i,'r')
99            self.assertTrue(x.equal(c))
100            self.A.restoreDenseColumnVec(i,'r', c)
101            self.assertFalse(c)
102            c = self.A.getDenseColumnVec(i,'w')
103            c.set(0)
104            self.A.restoreDenseColumnVec(i,'w', c)
105            self.assertFalse(c)
106            if i > 0:
107                c = self.A.getDenseColumnVec(i-1, mode='r')
108                self.assertEqual(c.norm(), 0)
109                self.A.restoreDenseColumnVec(i-1, V=c, mode='r')
110                self.assertFalse(c)
111        x.destroy()
112
113    def testCreateTranspose(self):
114        self._preallocate()
115        self._set_values()
116        self.A.assemble()
117        A = self.A
118        AT = PETSc.Mat().createTranspose(A)
119        x, y = A.createVecs()
120        xt, yt = AT.createVecs()
121        #
122        y.setRandom()
123        A.multTranspose(y, x)
124        y.copy(xt)
125        AT.mult(xt, yt)
126        self.assertTrue(yt.equal(x))
127        #
128        x.setRandom()
129        A.mult(x, y)
130        x.copy(yt)
131        AT.multTranspose(yt, xt)
132        self.assertTrue(xt.equal(y))
133        #
134        underlyingA = AT.getTransposeMat()
135        self.assertTrue(underlyingA.equal(A))
136
137    def _preallocate(self):
138        self.A.setPreallocationDense(None)
139
140    def _set_values(self):
141        COMM = self.COMM
142        GM, GN = self.GRID
143        BS = self.BSIZE or 1
144        rows, cols, vals = mkdata(COMM, GM, GN, BS)
145        if not self.BSIZE:
146            setvalues = self.A.setValues
147        else:
148            setvalues = self.A.setValuesBlocked
149        setvalues(rows, cols, vals)
150        return rows, cols, vals
151
152    def _chk_bs(self, A, bs):
153        self.assertEqual(A.getBlockSize(), bs or 1)
154
155    def _chk_array(self, A, r, c, v):
156        return  # XXX
157        vals = self.A.getValues(r, c)
158        vals.shape = v.shape
159        self.assertTrue(np.allclose(vals, v))
160
161
162# -- Dense ---------------------
163
164
165class BaseTestMatDense(BaseTestMatAnyDense, unittest.TestCase):
166    COMM = PETSc.COMM_WORLD
167    GRID = 0, 0
168    BSIZE = None
169
170
171# -- Seq Dense --
172
173
174class TestMatSeqDense(BaseTestMatDense):
175    COMM = PETSc.COMM_SELF
176    TYPE = PETSc.Mat.Type.SEQDENSE
177
178
179class TestMatSeqDense_G23(TestMatSeqDense):
180    GRID = 2, 3
181
182
183class TestMatSeqDense_G45(TestMatSeqDense):
184    GRID = 4, 5
185
186
187class TestMatSeqDense_G77(TestMatSeqDense):
188    GRID = 7, 7
189
190
191class TestMatSeqDense_G89(TestMatSeqDense):
192    GRID = 8, 9
193
194
195# -- MPI Dense --
196
197
198class TestMatMPIDense(BaseTestMatDense):
199    COMM = PETSc.COMM_WORLD
200    TYPE = PETSc.Mat.Type.MPIDENSE
201
202
203class TestMatMPIDense_G23(TestMatMPIDense):
204    GRID = 2, 3
205
206
207class TestMatMPIDense_G45(TestMatMPIDense):
208    GRID = 4, 5
209
210
211class TestMatMPIDense_G77(TestMatMPIDense):
212    GRID = 7, 7
213
214
215class TestMatMPIDense_G89(TestMatMPIDense):
216    GRID = 8, 9
217
218
219# -- Dense + Block ---------------
220
221
222class BaseTestMatDense_B(BaseTestMatAnyDense, unittest.TestCase):
223    COMM = PETSc.COMM_WORLD
224    GRID = 0, 0
225    BSIZE = 1
226
227    def _preallocate(self):
228        # self.A.setBlockSize(self.BSIZE)
229        self.A.setPreallocationDense(None)
230        # self.A.setBlockSize(self.BSIZE)
231        self._chk_bs(self.A, self.BSIZE)
232
233
234# -- Seq Dense + Block --
235
236
237class TestMatSeqDense_B(BaseTestMatDense_B):
238    COMM = PETSc.COMM_SELF
239    TYPE = PETSc.Mat.Type.SEQDENSE
240
241
242# bs = 1
243class TestMatSeqDense_B_G23(TestMatSeqDense_B):
244    GRID = 2, 3
245
246
247class TestMatSeqDense_B_G45(TestMatSeqDense_B):
248    GRID = 4, 5
249
250
251class TestMatSeqDense_B_G89(TestMatSeqDense_B):
252    GRID = 8, 9
253
254
255# bs = 2
256class TestMatSeqDense_B_G23_B2(TestMatSeqDense_B_G23):
257    BSIZE = 2
258
259
260class TestMatSeqDense_B_G45_B2(TestMatSeqDense_B_G45):
261    BSIZE = 2
262
263
264class TestMatSeqDense_B_G89_B2(TestMatSeqDense_B_G89):
265    BSIZE = 2
266
267
268# bs = 3
269class TestMatSeqDense_B_G23_B3(TestMatSeqDense_B_G23):
270    BSIZE = 3
271
272
273class TestMatSeqDense_B_G45_B3(TestMatSeqDense_B_G45):
274    BSIZE = 3
275
276
277class TestMatSeqDense_B_G89_B3(TestMatSeqDense_B_G89):
278    BSIZE = 3
279
280
281# bs = 4
282class TestMatSeqDense_B_G23_B4(TestMatSeqDense_B_G23):
283    BSIZE = 4
284
285
286class TestMatSeqDense_B_G45_B4(TestMatSeqDense_B_G45):
287    BSIZE = 4
288
289
290class TestMatSeqDense_B_G89_B4(TestMatSeqDense_B_G89):
291    BSIZE = 4
292
293
294# bs = 5
295class TestMatSeqDense_B_G23_B5(TestMatSeqDense_B_G23):
296    BSIZE = 5
297
298
299class TestMatSeqDense_B_G45_B5(TestMatSeqDense_B_G45):
300    BSIZE = 5
301
302
303class TestMatSeqDense_B_G89_B5(TestMatSeqDense_B_G89):
304    BSIZE = 5
305
306
307# -- MPI Dense + Block --
308
309
310class TestMatMPIDense_B(BaseTestMatDense_B):
311    COMM = PETSc.COMM_WORLD
312    TYPE = PETSc.Mat.Type.MPIDENSE
313
314
315# bs = 1
316class TestMatMPIDense_B_G23(TestMatMPIDense_B):
317    GRID = 2, 3
318
319
320class TestMatMPIDense_B_G45(TestMatMPIDense_B):
321    GRID = 4, 5
322
323
324class TestMatMPIDense_B_G77(TestMatMPIDense_B):
325    GRID = 7, 7
326
327
328class TestMatMPIDense_B_G89(TestMatMPIDense_B):
329    GRID = 8, 9
330
331
332# bs = 2
333class TestMatMPIDense_B_G23_B2(TestMatMPIDense_B_G23):
334    BSIZE = 2
335
336
337class TestMatMPIDense_B_G45_B2(TestMatMPIDense_B_G45):
338    BSIZE = 2
339
340
341class TestMatMPIDense_B_G77_B2(TestMatMPIDense_B_G77):
342    BSIZE = 2
343
344
345class TestMatMPIDense_B_G89_B2(TestMatMPIDense_B_G89):
346    BSIZE = 2
347
348
349# bs = 3
350class TestMatMPIDense_B_G23_B3(TestMatMPIDense_B_G23):
351    BSIZE = 3
352
353
354class TestMatMPIDense_B_G45_B3(TestMatMPIDense_B_G45):
355    BSIZE = 3
356
357
358class TestMatMPIDense_B_G77_B3(TestMatMPIDense_B_G77):
359    BSIZE = 3
360
361
362class TestMatMPIDense_B_G89_B3(TestMatMPIDense_B_G89):
363    BSIZE = 3
364
365
366# bs = 4
367class TestMatMPIDense_B_G23_B4(TestMatMPIDense_B_G23):
368    BSIZE = 4
369
370
371class TestMatMPIDense_B_G45_B4(TestMatMPIDense_B_G45):
372    BSIZE = 4
373
374
375class TestMatMPIDense_B_G77_B4(TestMatMPIDense_B_G77):
376    BSIZE = 4
377
378
379class TestMatMPIDense_B_G89_B4(TestMatMPIDense_B_G89):
380    BSIZE = 4
381
382
383# bs = 5
384class TestMatMPIDense_B_G23_B5(TestMatMPIDense_B_G23):
385    BSIZE = 5
386
387
388class TestMatMPIDense_B_G45_B5(TestMatMPIDense_B_G45):
389    BSIZE = 5
390
391
392class TestMatMPIDense_B_G77_B5(TestMatMPIDense_B_G77):
393    BSIZE = 5
394
395
396class TestMatMPIDense_B_G89_B5(TestMatMPIDense_B_G89):
397    BSIZE = 5
398
399
400# -----
401
402if __name__ == '__main__':
403    unittest.main()
404