xref: /petsc/src/binding/petsc4py/test/test_mat_aij.py (revision a69119a591a03a9d906b29c0a4e9802e4d7c9795)
1from petsc4py import PETSc
2import unittest
3
4import numpy as N
5import numpy as np
6
7def mkgraph(comm, m, n):
8    start = m*n * comm.rank
9    end   = start + m*n
10    idt = PETSc.IntType
11    rows = []
12    for I in range(start, end) :
13        rows.append([])
14        adj = rows[-1]
15        i = I//n; j = I - i*n
16        if i> 0  : J = I-n; adj.append(J)
17        if j> 0  : J = I-1; adj.append(J)
18        adj.append(I)
19        if j< n-1: J = I+1; adj.append(J)
20        if i< m-1: J = I+n; adj.append(J)
21    nods = N.array(range(start, end), dtype=idt)
22    xadj = N.array([0]*(len(rows)+1), dtype=idt)
23    xadj[0] = 0
24    xadj[1:] = N.cumsum([len(r) for r in rows], dtype=idt)
25    if not rows: adjy = N.array([],dtype=idt)
26    else:        adjy = N.concatenate(rows).astype(idt)
27    return nods, xadj, adjy
28
29
30class BaseTestMatAnyAIJ(object):
31
32    COMM  = PETSc.COMM_NULL
33    TYPE  = None
34    GRID  = 0, 0
35    BSIZE = None
36
37    def setUp(self):
38        COMM   = self.COMM
39        GM, GN = self.GRID
40        BS     = self.BSIZE
41        #
42        try:
43            rbs, cbs = BS
44            rbs = rbs or 1
45            cbs = cbs or 1
46        except (TypeError, ValueError):
47            rbs = cbs = BS or 1
48        sdt = dtype = PETSc.ScalarType
49        self.rows, self.xadj, self.adjy = mkgraph(COMM, GM, GN)
50        self.vals = N.array(range(1, 1 + len(self.adjy)*rbs*cbs), dtype=sdt)
51        self.vals.shape = (-1, rbs, cbs)
52        #
53        m, n = GM, GN
54        rowsz = (m*n*rbs, None)
55        colsz = (m*n*cbs, None)
56        A = self.A = PETSc.Mat().create(comm=COMM)
57        A.setType(self.TYPE)
58        A.setSizes([rowsz, colsz], BS)
59
60    def tearDown(self):
61        self.A.destroy()
62        self.A = None
63
64    def testSetPreallocNNZ(self):
65        nnz = [5, 2]
66        self.A.setPreallocationNNZ(nnz)
67        self._chk_bs(self.A, self.BSIZE)
68        opt = PETSc.Mat.Option.NEW_NONZERO_ALLOCATION_ERR
69        self.A.setOption(opt, True)
70        ai, aj, av = self._set_values()
71        self.A.assemble()
72        self._chk_aij(self.A, ai, aj)
73        opt = PETSc.Mat.Option.NEW_NONZERO_LOCATION_ERR
74        self.A.setOption(opt, True)
75        ai, aj, av = self._set_values_ijv()
76        self.A.assemble()
77        self._chk_aij(self.A, ai, aj)
78
79    def testSetPreallocNNZ_2(self):
80        _, ai, _, _ =self._get_aijv()
81        d_nnz = N.diff(ai)
82        nnz = [d_nnz, 3]
83        self.A.setPreallocationNNZ(nnz)
84        self._chk_bs(self.A, self.BSIZE)
85        opt = PETSc.Mat.Option.NEW_NONZERO_ALLOCATION_ERR
86        self.A.setOption(opt, True)
87        ai, aj, av = self._set_values()
88        self.A.assemble()
89        self._chk_aij(self.A, ai, aj)
90        opt = PETSc.Mat.Option.NEW_NONZERO_LOCATION_ERR
91        self.A.setOption(opt, True)
92        ai, aj, av =self._set_values_ijv()
93        self.A.assemble()
94        self._chk_aij(self.A, ai, aj)
95
96    def testSetPreallocCSR(self):
97        if 'is' in self.A.getType(): return # XXX
98        _, ai, aj, _ = self._get_aijv()
99        csr = [ai, aj]
100        self.A.setPreallocationCSR(csr)
101        self._chk_bs(self.A, self.BSIZE)
102        self._chk_aij(self.A, ai, aj)
103        opt = PETSc.Mat.Option.NEW_NONZERO_LOCATION_ERR
104        self.A.setOption(opt, True)
105        self._set_values()
106        self.A.assemble()
107        self._chk_aij(self.A, ai, aj)
108        self._set_values_ijv()
109        self.A.assemble()
110        self._chk_aij(self.A, ai, aj)
111
112    def testSetPreallocCSR_2(self):
113        if 'is' in self.A.getType(): return # XXX
114        _, ai, aj, av =self._get_aijv()
115        csr = [ai, aj, av]
116        self.A.setPreallocationCSR(csr)
117        self._chk_bs(self.A, self.BSIZE)
118        self._chk_aij(self.A, ai, aj)
119        opt = PETSc.Mat.Option.NEW_NONZERO_LOCATION_ERR
120        self.A.setOption(opt, True)
121        self._set_values()
122        self.A.assemble()
123        self._chk_aij(self.A, ai, aj)
124        self._set_values_ijv()
125        self.A.assemble()
126        self._chk_aij(self.A, ai, aj)
127
128    def testSetValues(self):
129        self._preallocate()
130        opt = PETSc.Mat.Option.NEW_NONZERO_ALLOCATION_ERR
131        self.A.setOption(opt, True)
132        ai, aj, av = self._set_values()
133        self.A.assemble()
134        self._chk_aij(self.A, ai, aj)
135        opt = PETSc.Mat.Option.NEW_NONZERO_LOCATION_ERR
136        self.A.setOption(opt, True)
137        ai, aj, av = self._set_values()
138        self.A.assemble()
139        self._chk_aij(self.A, ai, aj)
140
141    def testSetValuesIJV(self):
142        self._preallocate()
143        opt = PETSc.Mat.Option.NEW_NONZERO_ALLOCATION_ERR
144        self.A.setOption(opt, True)
145        ai, aj, av = self._set_values_ijv()
146        self.A.assemble()
147        self._chk_aij(self.A, ai, aj)
148        opt = PETSc.Mat.Option.NEW_NONZERO_LOCATION_ERR
149        self.A.setOption(opt, True)
150        ai, aj, av = self._set_values_ijv()
151        self.A.assemble()
152        self._chk_aij(self.A, ai, aj)
153
154    def testGetValuesCSR(self):
155        if 'is' in self.A.getType(): return # XXX
156        self._preallocate()
157        self._set_values_ijv()
158        A = self.A
159        A.assemble()
160        if 'sbaij' in A.getType():
161            opt = PETSc.Mat.Option.GETROW_UPPERTRIANGULAR
162            self.A.setOption(opt, True)
163        ai, aj, av = A.getValuesCSR()
164        rstart, rend = A.getOwnershipRange()
165        for row in range(rstart, rend):
166            cols, vals = A.getRow(row)
167            i = row - rstart
168            self.assertTrue(N.allclose(aj[ai[i]:ai[i+1]], cols))
169            self.assertTrue(N.allclose(av[ai[i]:ai[i+1]], vals))
170
171    def testConvertToSAME(self):
172        self._preallocate()
173        self._set_values_ijv()
174        A = self.A
175        A.assemble()
176        A.convert('same')
177
178    def testConvertToDENSE(self):
179        self._preallocate()
180        self._set_values_ijv()
181        A = self.A
182        A.assemble()
183        x, y = A.getVecs()
184        x.setRandom()
185        z = y.duplicate()
186        A.mult(x, y)
187        if A.type.endswith('sbaij'): return
188        B = PETSc.Mat()
189        A.convert('dense', B)  # initial
190        B.mult(x, z)
191        self.assertTrue(np.allclose(y.array, z.array))
192        A.convert('dense', B)  # reuse
193        B.mult(x, z)
194        self.assertTrue(np.allclose(y.array, z.array))
195        A.convert('dense')     # inplace
196        A.mult(x, z)
197        self.assertTrue(np.allclose(y.array, z.array))
198
199    def testConvertToAIJ(self):
200        self._preallocate()
201        self._set_values_ijv()
202        A = self.A
203        A.assemble()
204        x, y = A.getVecs()
205        x.setRandom()
206        z = y.duplicate()
207        A.mult(x, y)
208        if A.type.endswith('sbaij'): return
209        B = PETSc.Mat()
210        A.convert('aij', B)  # initial
211        B.mult(x, z)
212        self.assertTrue(np.allclose(y.array, z.array))
213        A.convert('aij', B)  # reuse
214        B.mult(x, z)
215        self.assertTrue(np.allclose(y.array, z.array))
216        A.convert('aij')     # inplace
217        A.mult(x, z)
218        self.assertTrue(np.allclose(y.array, z.array))
219
220    def testGetDiagonalBlock(self):
221        if 'is' in self.A.getType(): return # XXX
222        self._preallocate()
223        self._set_values_ijv()
224        self.A.assemble()
225        B = self.A.getDiagonalBlock()
226        self.assertEqual(self.A.getLocalSize(), B.getSize())
227        B.destroy()
228
229    def testInvertBlockDiagonal(self):
230        if 'is' in self.A.getType(): return # XXX
231        try:
232            _ = len(self.BSIZE)
233            return
234        except (TypeError, ValueError):
235            pass
236        self._preallocate()
237        rbs, cbs = self.A.getBlockSizes()
238        if rbs != cbs: return
239        self._set_values_ijv()
240        self.A.assemble()
241        self.A.shift(1000) # Make nonsingular
242        ibdiag = self.A.invertBlockDiagonal()
243        bs = self.A.getBlockSize()
244        m, _ = self.A.getLocalSize()
245        self.assertEqual(ibdiag.shape, (m//bs, bs, bs))
246        tmp = N.empty((m//bs, bs, bs), dtype=PETSc.ScalarType)
247        rstart, rend = self.A.getOwnershipRange()
248        s, e = rstart//bs, rend//bs
249        for i in range(s, e):
250            rows = cols = N.arange(i*bs,(i+1)*bs, dtype=PETSc.IntType)
251            vals = self.A.getValues(rows,cols)
252            tmp[i-s,:,:] = N.linalg.inv(vals)
253        self.assertTrue(N.allclose(ibdiag, tmp))
254
255    def testCreateSubMatrix(self):
256        if 'baij' in self.A.getType(): return # XXX
257        self._preallocate()
258        self._set_values_ijv()
259        self.A.assemble()
260        #
261        rank = self.A.getComm().getRank()
262        rs, re = self.A.getOwnershipRange()
263        cs, ce = self.A.getOwnershipRangeColumn()
264        rows = N.array(range(rs, re), dtype=PETSc.IntType)
265        cols = N.array(range(cs, ce), dtype=PETSc.IntType)
266        rows = PETSc.IS().createGeneral(rows, comm=self.A.getComm())
267        cols = PETSc.IS().createGeneral(cols, comm=self.A.getComm())
268        #
269        S = self.A.createSubMatrix(rows, None)
270        S.zeroEntries()
271        self.A.createSubMatrix(rows, None, S)
272        S.destroy()
273        #
274        S = self.A.createSubMatrix(rows, cols)
275        S.zeroEntries()
276        self.A.createSubMatrix(rows, cols, S)
277        S.destroy()
278
279    def testCreateSubMatrices(self):
280        if 'baij' in self.A.getType(): return # XXX
281        if 'is' in self.A.getType(): return # XXX
282        self._preallocate()
283        self._set_values_ijv()
284        self.A.assemble()
285        #
286        rs, re = self.A.getOwnershipRange()
287        cs, ce = self.A.getOwnershipRangeColumn()
288        rows = N.array(range(rs, re), dtype=PETSc.IntType)
289        cols = N.array(range(cs, ce), dtype=PETSc.IntType)
290        rows = PETSc.IS().createGeneral(rows, comm=self.A.getComm())
291        cols = PETSc.IS().createGeneral(cols, comm=self.A.getComm())
292        #
293        (S,) = self.A.createSubMatrices(rows, cols)
294        S.zeroEntries()
295        self.A.createSubMatrices(rows, cols, submats=[S])
296        S.destroy()
297        #
298        (S1,) = self.A.createSubMatrices([rows], [cols])
299        (S2,) = self.A.createSubMatrices([rows], [cols])
300        self.assertTrue(S1.equal(S2))
301        S2.zeroEntries()
302        self.A.createSubMatrices([rows], [cols], [S2])
303        self.assertTrue(S1.equal(S2))
304        S1.destroy()
305        S2.destroy()
306        #
307        if 'seq' not in self.A.getType(): return # XXX
308        S1, S2 = self.A.createSubMatrices([rows, rows], [cols, cols])
309        self.assertTrue(S1.equal(S2))
310        S1.zeroEntries()
311        S2.zeroEntries()
312        self.A.createSubMatrices([rows, rows], [cols, cols], [S1, S2])
313        self.assertTrue(S1.equal(S2))
314        S1.destroy()
315        S2.destroy()
316
317    def testGetRedundantMatrix(self):
318        if 'aijcrl' in self.A.getType(): return # duplicate not supported
319        if 'mpisbaij' in self.A.getType(): return # not working
320        if 'is' in self.A.getType(): return # XXX
321        self._preallocate()
322        self._set_values_ijv()
323        self.A.assemble()
324        #Test the most simple case
325        sizecommA = self.A.getComm().getSize()
326        Ared = self.A.getRedundantMatrix(sizecommA)
327        sizecommAred = Ared.getComm().getSize()
328        self.assertEqual(1, sizecommAred)
329        Ared.destroy()
330
331    def testCreateTranspose(self):
332        self._preallocate()
333        self._set_values_ijv()
334        self.A.assemble()
335        A = self.A
336        AT = PETSc.Mat().createTranspose(A)
337        x, y = A.createVecs()
338        xt, yt = AT.createVecs()
339        #
340        y.setRandom()
341        A.multTranspose(y, x)
342        y.copy(xt)
343        AT.mult(xt, yt)
344        self.assertTrue(yt.equal(x))
345        #
346        x.setRandom()
347        A.mult(x, y)
348        x.copy(yt)
349        AT.multTranspose(yt, xt)
350        self.assertTrue(xt.equal(y))
351
352    def _get_aijv(self):
353        return (self.rows, self.xadj, self.adjy, self.vals,)
354
355    def _preallocate(self):
356        self.A.setPreallocationNNZ([5, 2])
357
358    def _set_values(self):
359        import sys
360        if hasattr(sys, 'gettotalrefcount'):
361            return self._set_values_ijv()
362        # XXX Why the code below leak refs as a beast ???
363        row, ai, aj, av =self._get_aijv()
364        if not self.BSIZE:
365            setvalues = self.A.setValues
366        else:
367            setvalues = self.A.setValuesBlocked
368        for i, r in enumerate(row):
369            s, e = ai[i], ai[i+1]
370            setvalues(r, aj[s:e], av[s:e])
371        return ai, aj, av
372
373    def _set_values_ijv(self):
374        row, ai, aj, av =self._get_aijv()
375        if not self.BSIZE:
376            setvalues = self.A.setValuesIJV
377        else:
378            setvalues = self.A.setValuesBlockedIJV
379        setvalues(ai, aj, av, rowmap=row)
380        setvalues(ai, aj, av, rowmap=None)
381        return ai, aj, av
382
383    def _chk_bs(self, A, bs):
384        self.assertEqual(A.getBlockSize(), bs or 1)
385
386    def _chk_bsizes(self, A, bsizes):
387        try:
388            rbs, cbs = bsizes
389        except (TypeError, ValueError):
390            rbs = cbs = bsizes
391        self.assertEqual(A.getBlockSizes(), (rbs, cbs))
392
393    def _chk_aij(self, A, i, j):
394        compressed = bool(self.BSIZE)
395        ai, aj = A.getRowIJ(compressed=compressed)
396        if ai is not None and aj is not None:
397            self.assertTrue(N.all(i==ai))
398            self.assertTrue(N.all(j==aj))
399        ai, aj = A.getColumnIJ(compressed=compressed)
400        if ai is not None and aj is not None:
401            self.assertTrue(N.all(i==ai))
402            self.assertTrue(N.all(j==aj))
403
404# -- AIJ ---------------------
405
406class BaseTestMatAIJ(BaseTestMatAnyAIJ, unittest.TestCase):
407    COMM  = PETSc.COMM_WORLD
408    TYPE  = PETSc.Mat.Type.AIJ
409    GRID  = 0, 0
410    BSIZE = None
411
412# -- Seq AIJ --
413
414class TestMatSeqAIJ(BaseTestMatAIJ):
415    COMM = PETSc.COMM_SELF
416    TYPE = PETSc.Mat.Type.SEQAIJ
417class TestMatSeqAIJ_G23(TestMatSeqAIJ):
418    GRID  = 2, 3
419class TestMatSeqAIJ_G45(TestMatSeqAIJ):
420    GRID  = 4, 5
421class TestMatSeqAIJ_G89(TestMatSeqAIJ):
422    GRID  = 8, 9
423
424# -- MPI AIJ --
425
426class TestMatMPIAIJ(BaseTestMatAIJ):
427    COMM = PETSc.COMM_WORLD
428    TYPE = PETSc.Mat.Type.MPIAIJ
429class TestMatMPIAIJ_G23(TestMatMPIAIJ):
430    GRID  = 2, 3
431class TestMatMPIAIJ_G45(TestMatMPIAIJ):
432    GRID  = 4, 5
433class TestMatMPIAIJ_G89(TestMatMPIAIJ):
434    GRID  = 8, 9
435
436
437# -- Block AIJ ---------------
438
439class BaseTestMatBAIJ(BaseTestMatAnyAIJ, unittest.TestCase):
440    COMM  = PETSc.COMM_WORLD
441    TYPE  = PETSc.Mat.Type.BAIJ
442    GRID  = 0, 0
443    BSIZE = 1
444
445# -- Seq Block AIJ --
446
447class TestMatSeqBAIJ(BaseTestMatBAIJ):
448    COMM = PETSc.COMM_SELF
449    TYPE = PETSc.Mat.Type.SEQBAIJ
450# bs = 1
451class TestMatSeqBAIJ_G23(TestMatSeqBAIJ):
452    GRID  = 2, 3
453class TestMatSeqBAIJ_G45(TestMatSeqBAIJ):
454    GRID  = 4, 5
455class TestMatSeqBAIJ_G89(TestMatSeqBAIJ):
456    GRID  = 8, 9
457# bs = 2
458class TestMatSeqBAIJ_G23_B2(TestMatSeqBAIJ_G23):
459    BSIZE = 2
460class TestMatSeqBAIJ_G45_B2(TestMatSeqBAIJ_G45):
461    BSIZE = 2
462class TestMatSeqBAIJ_G89_B2(TestMatSeqBAIJ_G89):
463    BSIZE = 2
464# bs = 3
465class TestMatSeqBAIJ_G23_B3(TestMatSeqBAIJ_G23):
466    BSIZE = 3
467class TestMatSeqBAIJ_G45_B3(TestMatSeqBAIJ_G45):
468    BSIZE = 3
469class TestMatSeqBAIJ_G89_B3(TestMatSeqBAIJ_G89):
470    BSIZE = 3
471# bs = 4
472class TestMatSeqBAIJ_G23_B4(TestMatSeqBAIJ_G23):
473    BSIZE = 4
474class TestMatSeqBAIJ_G45_B4(TestMatSeqBAIJ_G45):
475    BSIZE = 4
476class TestMatSeqBAIJ_G89_B4(TestMatSeqBAIJ_G89):
477    BSIZE = 4
478# bs = 5
479class TestMatSeqBAIJ_G23_B5(TestMatSeqBAIJ_G23):
480    BSIZE = 5
481class TestMatSeqBAIJ_G45_B5(TestMatSeqBAIJ_G45):
482    BSIZE = 5
483class TestMatSeqBAIJ_G89_B5(TestMatSeqBAIJ_G89):
484    BSIZE = 5
485
486
487# -- MPI Block AIJ --
488
489class TestMatMPIBAIJ(BaseTestMatBAIJ):
490    COMM = PETSc.COMM_WORLD
491    TYPE = PETSc.Mat.Type.MPIBAIJ
492# bs = 1
493class TestMatMPIBAIJ_G23(TestMatMPIBAIJ):
494    GRID  = 2, 3
495class TestMatMPIBAIJ_G45(TestMatMPIBAIJ):
496    GRID  = 4, 5
497class TestMatMPIBAIJ_G89(TestMatMPIBAIJ):
498    GRID  = 8, 9
499# bs = 2
500class TestMatMPIBAIJ_G23_B2(TestMatMPIBAIJ_G23):
501    BSIZE = 2
502class TestMatMPIBAIJ_G45_B2(TestMatMPIBAIJ_G45):
503    BSIZE = 2
504class TestMatMPIBAIJ_G89_B2(TestMatMPIBAIJ_G89):
505    BSIZE = 2
506# bs = 3
507class TestMatMPIBAIJ_G23_B3(TestMatMPIBAIJ_G23):
508    BSIZE = 3
509class TestMatMPIBAIJ_G45_B3(TestMatMPIBAIJ_G45):
510    BSIZE = 3
511class TestMatMPIBAIJ_G89_B3(TestMatMPIBAIJ_G89):
512    BSIZE = 3
513# bs = 4
514class TestMatMPIBAIJ_G23_B4(TestMatMPIBAIJ_G23):
515    BSIZE = 4
516class TestMatMPIBAIJ_G45_B4(TestMatMPIBAIJ_G45):
517    BSIZE = 4
518class TestMatMPIBAIJ_G89_B4(TestMatMPIBAIJ_G89):
519    BSIZE = 4
520# bs = 5
521class TestMatMPIBAIJ_G23_B5(TestMatMPIBAIJ_G23):
522    BSIZE = 5
523class TestMatMPIBAIJ_G45_B5(TestMatMPIBAIJ_G45):
524    BSIZE = 5
525class TestMatMPIBAIJ_G89_B5(TestMatMPIBAIJ_G89):
526    BSIZE = 5
527
528# -- SymmBlock AIJ ---------------
529
530class BaseTestMatSBAIJ(BaseTestMatAnyAIJ, unittest.TestCase):
531    COMM  = PETSc.COMM_WORLD
532    TYPE  = PETSc.Mat.Type.SBAIJ
533    GRID  = 0, 0
534    BSIZE = 1
535    def testInvertBlockDiagonal(self): pass
536    def _chk_aij(self, A, i, j):
537        ai, aj = A.getRowIJ(compressed=True)
538        if ai is not None and aj is not None:
539            if 0: # XXX Implement
540                self.assertTrue(N.all(i==ai))
541                self.assertTrue(N.all(j==aj))
542        ai, aj = A.getColumnIJ(compressed=True)
543        if ai is not None and aj is not None:
544            if 0: # XXX Implement
545                self.assertTrue(N.all(i==ai))
546                self.assertTrue(N.all(j==aj))
547
548# -- Seq SymmBlock AIJ --
549
550class TestMatSeqSBAIJ(BaseTestMatSBAIJ):
551    COMM = PETSc.COMM_SELF
552    TYPE = PETSc.Mat.Type.SEQSBAIJ
553# bs = 1
554class TestMatSeqSBAIJ_G23(TestMatSeqSBAIJ):
555    GRID  = 2, 3
556class TestMatSeqSBAIJ_G45(TestMatSeqSBAIJ):
557    GRID  = 4, 5
558class TestMatSeqSBAIJ_G89(TestMatSeqSBAIJ):
559    GRID  = 8, 9
560# bs = 2
561class TestMatSeqSBAIJ_G23_B2(TestMatSeqSBAIJ_G23):
562    BSIZE = 2
563class TestMatSeqSBAIJ_G45_B2(TestMatSeqSBAIJ_G45):
564    BSIZE = 2
565class TestMatSeqSBAIJ_G89_B2(TestMatSeqSBAIJ_G89):
566    BSIZE = 2
567# bs = 3
568class TestMatSeqSBAIJ_G23_B3(TestMatSeqSBAIJ_G23):
569    BSIZE = 3
570class TestMatSeqSBAIJ_G45_B3(TestMatSeqSBAIJ_G45):
571    BSIZE = 3
572class TestMatSeqSBAIJ_G89_B3(TestMatSeqSBAIJ_G89):
573    BSIZE = 3
574# bs = 4
575class TestMatSeqSBAIJ_G23_B4(TestMatSeqSBAIJ_G23):
576    BSIZE = 4
577class TestMatSeqSBAIJ_G45_B4(TestMatSeqSBAIJ_G45):
578    BSIZE = 4
579class TestMatSeqSBAIJ_G89_B4(TestMatSeqSBAIJ_G89):
580    BSIZE = 4
581# bs = 5
582class TestMatSeqSBAIJ_G23_B5(TestMatSeqSBAIJ_G23):
583    BSIZE = 5
584class TestMatSeqSBAIJ_G45_B5(TestMatSeqSBAIJ_G45):
585    BSIZE = 5
586class TestMatSeqSBAIJ_G89_B5(TestMatSeqSBAIJ_G89):
587    BSIZE = 5
588
589
590# -- MPI SymmBlock AIJ --
591
592class TestMatMPISBAIJ(BaseTestMatSBAIJ):
593    COMM = PETSc.COMM_WORLD
594    TYPE = PETSc.Mat.Type.MPISBAIJ
595# bs = 1
596class TestMatMPISBAIJ_G23(TestMatMPISBAIJ):
597    GRID  = 2, 3
598class TestMatMPISBAIJ_G45(TestMatMPISBAIJ):
599    GRID  = 4, 5
600class TestMatMPISBAIJ_G89(TestMatMPISBAIJ):
601    GRID  = 8, 9
602# bs = 2
603class TestMatMPISBAIJ_G23_B2(TestMatMPISBAIJ_G23):
604    BSIZE = 2
605class TestMatMPISBAIJ_G45_B2(TestMatMPISBAIJ_G45):
606    BSIZE = 2
607class TestMatMPISBAIJ_G89_B2(TestMatMPISBAIJ_G89):
608    BSIZE = 2
609# bs = 3
610class TestMatMPISBAIJ_G23_B3(TestMatMPISBAIJ_G23):
611    BSIZE = 3
612class TestMatMPISBAIJ_G45_B3(TestMatMPISBAIJ_G45):
613    BSIZE = 3
614class TestMatMPISBAIJ_G89_B3(TestMatMPISBAIJ_G89):
615    BSIZE = 3
616# bs = 4
617class TestMatMPISBAIJ_G23_B4(TestMatMPISBAIJ_G23):
618    BSIZE = 4
619class TestMatMPISBAIJ_G45_B4(TestMatMPISBAIJ_G45):
620    BSIZE = 4
621class TestMatMPISBAIJ_G89_B4(TestMatMPISBAIJ_G89):
622    BSIZE = 4
623# bs = 5
624class TestMatMPISBAIJ_G23_B5(TestMatMPISBAIJ_G23):
625    BSIZE = 5
626class TestMatMPISBAIJ_G45_B5(TestMatMPISBAIJ_G45):
627    BSIZE = 5
628class TestMatMPISBAIJ_G89_B5(TestMatMPISBAIJ_G89):
629    BSIZE = 5
630
631# -- AIJ + Block ---------------
632
633class BaseTestMatAIJ_B(BaseTestMatAnyAIJ, unittest.TestCase):
634    COMM  = PETSc.COMM_WORLD
635    TYPE  = PETSc.Mat.Type.AIJ
636    GRID  = 0, 0
637    BSIZE = 1
638
639    def testSetPreallocNNZ(self):pass
640    def testSetPreallocNNZ_2(self):pass
641    def testSetPreallocCSR(self):pass
642    def testSetPreallocCSR_2(self):pass
643    def testSetValues(self):
644        self._preallocate()
645        opt = PETSc.Mat.Option.NEW_NONZERO_ALLOCATION_ERR
646        self.A.setOption(opt, True)
647        ai, aj, av = self._set_values()
648        self.A.assemble()
649        self._chk_aij(self.A, ai, aj)
650        opt = PETSc.Mat.Option.NEW_NONZERO_LOCATION_ERR
651        self.A.setOption(opt, True)
652        ai, aj, av = self._set_values()
653        self.A.assemble()
654        self._chk_aij(self.A, ai, aj)
655    def testSetValuesIJV(self):
656        self._preallocate()
657        opt = PETSc.Mat.Option.NEW_NONZERO_ALLOCATION_ERR
658        self.A.setOption(opt, True)
659        ai, aj, av = self._set_values_ijv()
660        self.A.assemble()
661        self._chk_aij(self.A, ai, aj)
662        opt = PETSc.Mat.Option.NEW_NONZERO_LOCATION_ERR
663        self.A.setOption(opt, True)
664        ai, aj, av = self._set_values_ijv()
665        self.A.assemble()
666        self._chk_aij(self.A, ai, aj)
667    def _preallocate(self):
668        self.A.setPreallocationNNZ([5*self.BSIZE, 3*self.BSIZE])
669        self._chk_bs(self.A, self.BSIZE)
670    def _chk_aij(self, A, i, j):
671        bs = self.BSIZE or 1
672        ai, aj = A.getRowIJ()
673        if ai is not None and aj is not None:  ## XXX map and check !!
674            #self.assertTrue(N.all(i==ai))
675            #self.assertTrue(N.all(j==aj))
676            pass
677        ai, aj = A.getColumnIJ(compressed=bool(self.BSIZE))
678        if ai is not None and aj is not None: ## XXX map and check !!
679            #self.assertTrue(N.all(i==ai))
680            #self.assertTrue(N.all(j==aj))
681            pass
682
683# -- Seq AIJ + Block --
684
685class TestMatSeqAIJ_B(BaseTestMatAIJ_B):
686    COMM = PETSc.COMM_SELF
687    TYPE = PETSc.Mat.Type.SEQAIJ
688# bs = 1
689class TestMatSeqAIJ_B_G23(TestMatSeqAIJ_B):
690    GRID  = 2, 3
691class TestMatSeqAIJ_B_G45(TestMatSeqAIJ_B):
692    GRID  = 4, 5
693class TestMatSeqAIJ_B_G89(TestMatSeqAIJ_B):
694    GRID  = 8, 9
695# bs = 2
696class TestMatSeqAIJ_B_G23_B2(TestMatSeqAIJ_B_G23):
697    BSIZE = 2
698class TestMatSeqAIJ_B_G45_B2(TestMatSeqAIJ_B_G45):
699    BSIZE = 2
700class TestMatSeqAIJ_B_G89_B2(TestMatSeqAIJ_B_G89):
701    BSIZE = 2
702# bs = 3
703class TestMatSeqAIJ_B_G23_B3(TestMatSeqAIJ_B_G23):
704    BSIZE = 3
705class TestMatSeqAIJ_B_G45_B3(TestMatSeqAIJ_B_G45):
706    BSIZE = 3
707class TestMatSeqAIJ_B_G89_B3(TestMatSeqAIJ_B_G89):
708    BSIZE = 3
709# bs = 4
710class TestMatSeqAIJ_B_G23_B4(TestMatSeqAIJ_B_G23):
711    BSIZE = 4
712class TestMatSeqAIJ_B_G45_B4(TestMatSeqAIJ_B_G45):
713    BSIZE = 4
714class TestMatSeqAIJ_B_G89_B4(TestMatSeqAIJ_B_G89):
715    BSIZE = 4
716# bs = 5
717class TestMatSeqAIJ_B_G23_B5(TestMatSeqAIJ_B_G23):
718    BSIZE = 5
719class TestMatSeqAIJ_B_G45_B5(TestMatSeqAIJ_B_G45):
720    BSIZE = 5
721class TestMatSeqAIJ_B_G89_B5(TestMatSeqAIJ_B_G89):
722    BSIZE = 5
723
724
725# -- MPI AIJ + Block --
726
727class TestMatMPIAIJ_B(BaseTestMatAIJ_B):
728    COMM = PETSc.COMM_WORLD
729    TYPE = PETSc.Mat.Type.MPIAIJ
730# bs = 1
731class TestMatMPIAIJ_B_G23(TestMatMPIAIJ_B):
732    GRID  = 2, 3
733class TestMatMPIAIJ_B_G45(TestMatMPIAIJ_B):
734    GRID  = 4, 5
735class TestMatMPIAIJ_B_G89(TestMatMPIAIJ_B):
736    GRID  = 8, 9
737# bs = 2
738class TestMatMPIAIJ_B_G23_B2(TestMatMPIAIJ_B_G23):
739    BSIZE = 2
740class TestMatMPIAIJ_B_G45_B2(TestMatMPIAIJ_B_G45):
741    BSIZE = 2
742class TestMatMPIAIJ_B_G89_B2(TestMatMPIAIJ_B_G89):
743    BSIZE = 2
744# bs = 3
745class TestMatMPIAIJ_B_G23_B3(TestMatMPIAIJ_B_G23):
746    BSIZE = 3
747class TestMatMPIAIJ_B_G45_B3(TestMatMPIAIJ_B_G45):
748    BSIZE = 3
749class TestMatMPIAIJ_B_G89_B3(TestMatMPIAIJ_B_G89):
750    BSIZE = 3
751# bs = 4
752class TestMatMPIAIJ_B_G23_B4(TestMatMPIAIJ_B_G23):
753    BSIZE = 4
754class TestMatMPIAIJ_B_G45_B4(TestMatMPIAIJ_B_G45):
755    BSIZE = 4
756class TestMatMPIAIJ_B_G89_B4(TestMatMPIAIJ_B_G89):
757    BSIZE = 4
758# bs = 5
759class TestMatMPIAIJ_B_G23_B5(TestMatMPIAIJ_B_G23):
760    BSIZE = 5
761class TestMatMPIAIJ_B_G45_B5(TestMatMPIAIJ_B_G45):
762    BSIZE = 5
763class TestMatMPIAIJ_B_G89_B5(TestMatMPIAIJ_B_G89):
764    BSIZE = 5
765
766# -- Non-square blocks --
767class BaseTestMatAIJ_B(BaseTestMatAnyAIJ, unittest.TestCase):
768    COMM  = PETSc.COMM_WORLD
769    TYPE  = PETSc.Mat.Type.AIJ
770    GRID  = 0, 0
771    BSIZE = 4, 2
772
773    def _preallocate(self):
774        try:
775            rbs, cbs = self.BSIZE
776        except (TypeError, ValueError):
777            rbs = cbs = self.BSIZE
778        self.A.setPreallocationNNZ([5*rbs, 3*cbs])
779        self._chk_bsizes(self.A, self.BSIZE)
780    def testSetPreallocNNZ(self):pass
781    def testSetPreallocNNZ_2(self):pass
782    def testSetPreallocCSR(self):pass
783    def testSetPreallocCSR_2(self):pass
784    def testSetValues(self):
785        self._preallocate()
786        opt = PETSc.Mat.Option.NEW_NONZERO_ALLOCATION_ERR
787        self.A.setOption(opt, True)
788        ai, aj, av = self._set_values()
789        self.A.assemble()
790        self._chk_aij(self.A, ai, aj)
791        opt = PETSc.Mat.Option.NEW_NONZERO_LOCATION_ERR
792        self.A.setOption(opt, True)
793        ai, aj, av = self._set_values()
794        self.A.assemble()
795        self._chk_aij(self.A, ai, aj)
796    def testSetValuesIJV(self):
797        self._preallocate()
798        opt = PETSc.Mat.Option.NEW_NONZERO_ALLOCATION_ERR
799        self.A.setOption(opt, True)
800        ai, aj, av = self._set_values_ijv()
801        self.A.assemble()
802        self._chk_aij(self.A, ai, aj)
803        opt = PETSc.Mat.Option.NEW_NONZERO_LOCATION_ERR
804        self.A.setOption(opt, True)
805        ai, aj, av = self._set_values_ijv()
806        self.A.assemble()
807        self._chk_aij(self.A, ai, aj)
808    def _chk_aij(self, A, i, j):
809        bs = self.BSIZE or 1
810        ai, aj = A.getRowIJ()
811        if ai is not None and aj is not None:  ## XXX map and check !!
812            #self.assertTrue(N.all(i==ai))
813            #self.assertTrue(N.all(j==aj))
814            pass
815        ai, aj = A.getColumnIJ()
816        if ai is not None and aj is not None: ## XXX map and check !!
817            #self.assertTrue(N.all(i==ai))
818            #self.assertTrue(N.all(j==aj))
819            pass
820
821# -- AIJCRL ---------------------
822
823class BaseTestMatAIJCRL(BaseTestMatAIJ, unittest.TestCase):
824    TYPE  = PETSc.Mat.Type.AIJCRL
825
826# -- Seq AIJCRL --
827
828class TestMatSeqAIJCRL(BaseTestMatAIJCRL):
829    COMM = PETSc.COMM_SELF
830    TYPE = PETSc.Mat.Type.SEQAIJCRL
831class TestMatSeqAIJCRL_G23(TestMatSeqAIJCRL):
832    GRID  = 2, 3
833class TestMatSeqAIJCRL_G45(TestMatSeqAIJCRL):
834    GRID  = 4, 5
835class TestMatSeqAIJCRL_G89(TestMatSeqAIJCRL):
836    GRID  = 8, 9
837
838# -- MPI AIJCRL --
839
840class TestMatMPIAIJCRL(BaseTestMatAIJCRL):
841    COMM = PETSc.COMM_WORLD
842    TYPE = PETSc.Mat.Type.MPIAIJCRL
843class TestMatMPIAIJCRL_G23(TestMatMPIAIJCRL):
844    GRID  = 2, 3
845class TestMatMPIAIJCRL_G45(TestMatMPIAIJCRL):
846    GRID  = 4, 5
847class TestMatMPIAIJCRL_G89(TestMatMPIAIJCRL):
848    GRID  = 8, 9
849
850# -- AIJCRL + Block -------------
851
852class BaseTestMatAIJCRL_B(BaseTestMatAIJ_B, unittest.TestCase):
853    TYPE  = PETSc.Mat.Type.AIJCRL
854
855# -- Seq AIJCRL + Block --
856
857class TestMatSeqAIJCRL_B(BaseTestMatAIJCRL_B):
858    COMM = PETSc.COMM_SELF
859    TYPE = PETSc.Mat.Type.SEQAIJCRL
860# bs = 1
861class TestMatSeqAIJCRL_B_G23(TestMatSeqAIJCRL_B):
862    GRID  = 2, 3
863class TestMatSeqAIJCRL_B_G45(TestMatSeqAIJCRL_B):
864    GRID  = 4, 5
865class TestMatSeqAIJCRL_B_G89(TestMatSeqAIJCRL_B):
866    GRID  = 8, 9
867# bs = 2
868class TestMatSeqAIJCRL_B_G23_B2(TestMatSeqAIJCRL_B_G23):
869    BSIZE = 2
870class TestMatSeqAIJCRL_B_G45_B2(TestMatSeqAIJCRL_B_G45):
871    BSIZE = 2
872class TestMatSeqAIJCRL_B_G89_B2(TestMatSeqAIJCRL_B_G89):
873    BSIZE = 2
874# bs = 3
875class TestMatSeqAIJCRL_B_G23_B3(TestMatSeqAIJCRL_B_G23):
876    BSIZE = 3
877class TestMatSeqAIJCRL_B_G45_B3(TestMatSeqAIJCRL_B_G45):
878    BSIZE = 3
879class TestMatSeqAIJCRL_B_G89_B3(TestMatSeqAIJCRL_B_G89):
880    BSIZE = 3
881# bs = 4
882class TestMatSeqAIJCRL_B_G23_B4(TestMatSeqAIJCRL_B_G23):
883    BSIZE = 4
884class TestMatSeqAIJCRL_B_G45_B4(TestMatSeqAIJCRL_B_G45):
885    BSIZE = 4
886class TestMatSeqAIJCRL_B_G89_B4(TestMatSeqAIJCRL_B_G89):
887    BSIZE = 4
888# bs = 5
889class TestMatSeqAIJCRL_B_G23_B5(TestMatSeqAIJCRL_B_G23):
890    BSIZE = 5
891class TestMatSeqAIJCRL_B_G45_B5(TestMatSeqAIJCRL_B_G45):
892    BSIZE = 5
893class TestMatSeqAIJCRL_B_G89_B5(TestMatSeqAIJCRL_B_G89):
894    BSIZE = 5
895
896
897# -- MPI AIJCRL + Block --
898
899class TestMatMPIAIJCRL_B(BaseTestMatAIJCRL_B):
900    COMM = PETSc.COMM_WORLD
901    TYPE = PETSc.Mat.Type.MPIAIJCRL
902# bs = 1
903class TestMatMPIAIJCRL_B_G23(TestMatMPIAIJCRL_B):
904    GRID  = 2, 3
905class TestMatMPIAIJCRL_B_G45(TestMatMPIAIJCRL_B):
906    GRID  = 4, 5
907class TestMatMPIAIJCRL_B_G89(TestMatMPIAIJCRL_B):
908    GRID  = 8, 9
909# bs = 2
910class TestMatMPIAIJCRL_B_G23_B2(TestMatMPIAIJCRL_B_G23):
911    BSIZE = 2
912class TestMatMPIAIJCRL_B_G45_B2(TestMatMPIAIJCRL_B_G45):
913    BSIZE = 2
914class TestMatMPIAIJCRL_B_G89_B2(TestMatMPIAIJCRL_B_G89):
915    BSIZE = 2
916# bs = 3
917class TestMatMPIAIJCRL_B_G23_B3(TestMatMPIAIJCRL_B_G23):
918    BSIZE = 3
919class TestMatMPIAIJCRL_B_G45_B3(TestMatMPIAIJCRL_B_G45):
920    BSIZE = 3
921class TestMatMPIAIJCRL_B_G89_B3(TestMatMPIAIJCRL_B_G89):
922    BSIZE = 3
923# bs = 4
924class TestMatMPIAIJCRL_B_G23_B4(TestMatMPIAIJCRL_B_G23):
925    BSIZE = 4
926class TestMatMPIAIJCRL_B_G45_B4(TestMatMPIAIJCRL_B_G45):
927    BSIZE = 4
928class TestMatMPIAIJCRL_B_G89_B4(TestMatMPIAIJCRL_B_G89):
929    BSIZE = 4
930# bs = 5
931class TestMatMPIAIJCRL_B_G23_B5(TestMatMPIAIJCRL_B_G23):
932    BSIZE = 5
933class TestMatMPIAIJCRL_B_G45_B5(TestMatMPIAIJCRL_B_G45):
934    BSIZE = 5
935class TestMatMPIAIJCRL_B_G89_B5(TestMatMPIAIJCRL_B_G89):
936    BSIZE = 5
937
938# -- MATIS --
939
940class TestMatIS(BaseTestMatAIJ):
941    COMM = PETSc.COMM_WORLD
942    TYPE = PETSc.Mat.Type.IS
943class TestMatIS_G23(TestMatIS):
944    GRID  = 2, 3
945class TestMatIS_G45(TestMatIS):
946    GRID  = 4, 5
947class TestMatIS_G89(TestMatIS):
948    GRID  = 8, 9
949
950# -----
951
952
953
954if __name__ == '__main__':
955    unittest.main()
956