1from petsc4py import PETSc
2import unittest
3
4# --------------------------------------------------------------------
5
6class BaseTestDA(object):
7
8    COMM = PETSc.COMM_WORLD
9    SIZES = None
10    BOUNDARY = None
11    DOF = 1
12    STENCIL = PETSc.DMDA.StencilType.STAR
13    SWIDTH = 1
14
15    def setUp(self):
16        self.da = PETSc.DMDA().create(dim=len(self.SIZES),
17                                      dof=self.DOF,
18                                      sizes=self.SIZES,
19                                      boundary_type=self.BOUNDARY,
20                                      stencil_type=self.STENCIL,
21                                      stencil_width=self.SWIDTH,
22                                      comm=self.COMM)
23
24    def tearDown(self):
25        self.da = None
26
27    def testGetInfo(self):
28        dim = self.da.getDim()
29        dof = self.da.getDof()
30        sizes = self.da.getSizes()
31        psizes = self.da.getProcSizes()
32        boundary = self.da.getBoundaryType()
33        stencil_type = self.da.getStencilType()
34        stencil_width = self.da.getStencilWidth()
35        self.assertEqual(dim, len(self.SIZES))
36        self.assertEqual(dof, self.DOF)
37        self.assertEqual(sizes, tuple(self.SIZES))
38        self.assertEqual(boundary, self.BOUNDARY or (0,)*dim)
39        self.assertEqual(stencil_type, self.STENCIL)
40        self.assertEqual(stencil_width, self.SWIDTH)
41
42    def testRangesCorners(self):
43        dim = self.da.getDim()
44        ranges = self.da.getRanges()
45        starts, lsizes  = self.da.getCorners()
46        self.assertEqual(dim, len(ranges))
47        self.assertEqual(dim, len(starts))
48        self.assertEqual(dim, len(lsizes))
49        for i in range(dim):
50            s, e = ranges[i]
51            self.assertEqual(s, starts[i])
52            self.assertEqual(e-s, lsizes[i])
53
54    def testGhostRangesCorners(self):
55        dim = self.da.getDim()
56        ranges = self.da.getGhostRanges()
57        starts, lsizes  = self.da.getGhostCorners()
58        self.assertEqual(dim, len(ranges))
59        self.assertEqual(dim, len(starts))
60        self.assertEqual(dim, len(lsizes))
61        for i in range(dim):
62            s, e = ranges[i]
63            self.assertEqual(s, starts[i])
64            self.assertEqual(e-s, lsizes[i])
65
66    def testOwnershipRanges(self):
67        dim = self.da.getDim()
68        ownership_ranges = self.da.getOwnershipRanges()
69        procsizes = self.da.getProcSizes()
70        self.assertEqual(len(procsizes), len(ownership_ranges))
71        for i,m in enumerate(procsizes):
72            self.assertEqual(m, len(ownership_ranges[i]))
73
74    def testFieldName(self):
75        for i in range(self.da.getDof()):
76            self.da.setFieldName(i, "field%d" % i)
77        for i in range(self.da.getDof()):
78            name = self.da.getFieldName(i)
79            self.assertEqual(name, "field%d" % i)
80
81    def testCoordinates(self):
82        self.da.setUniformCoordinates(0,1,0,1,0,1)
83        #
84        c = self.da.getCoordinates()
85        self.da.setCoordinates(c)
86        c.destroy()
87        cda = self.da.getCoordinateDM()
88        cda.destroy()
89        #
90        c = self.da.getCoordinates()
91        self.da.setCoordinates(c)
92        c.destroy()
93        gc = self.da.getCoordinatesLocal()
94        gc.destroy()
95
96    def testCreateVecMat(self):
97        vn = self.da.createNaturalVec()
98        vg = self.da.createGlobalVec()
99        vl = self.da.createLocalVec()
100        mat = self.da.createMat()
101        self.assertTrue(mat.getType() in ('aij', 'seqaij', 'mpiaij'))
102        vn.set(1.0)
103        self.da.naturalToGlobal(vn,vg)
104        self.assertEqual(vg.max()[1], 1.0)
105        self.assertEqual(vg.min()[1], 1.0)
106        self.da.globalToLocal(vg,vl)
107        self.assertEqual(vl.max()[1], 1.0)
108        self.assertTrue (vl.min()[1] in (1.0, 0.0))
109        vn.set(0.0)
110        self.da.globalToNatural(vg,vn)
111        self.assertEqual(vn.max()[1], 1.0)
112        self.assertEqual(vn.min()[1], 1.0)
113        vl2 = self.da.createLocalVec()
114        self.da.localToLocal(vl,vl2)
115        self.assertEqual(vl2.max()[1], 1.0)
116        self.assertTrue (vl2.min()[1] in (1.0, 0.0))
117        NONE = PETSc.DM.BoundaryType.NONE
118        s = self.da.stencil_width
119        btype = self.da.boundary_type
120        psize = self.da.proc_sizes
121        for b, p in zip(btype, psize):
122            if b != NONE and p == 1: return
123        vg2 = self.da.createGlobalVec()
124        self.da.localToGlobal(vl2,vg2)
125
126    def testGetVec(self):
127        vg = self.da.getGlobalVec()
128        vl = self.da.getLocalVec()
129        try:
130            vg.set(1.0)
131            self.assertEqual(vg.max()[1], 1.0)
132            self.assertEqual(vg.min()[1], 1.0)
133            self.da.globalToLocal(vg,vl)
134            self.assertEqual(vl.max()[1], 1.0)
135            self.assertTrue (vl.min()[1] in (1.0, 0.0))
136            vl.set(2.0)
137            NONE = PETSc.DM.BoundaryType.NONE
138            s = self.da.stencil_width
139            btype = self.da.boundary_type
140            psize = self.da.proc_sizes
141            for b, p in zip(btype, psize):
142                if b != NONE and p == 1: return
143            self.da.localToGlobal(vl,vg)
144            self.assertEqual(vg.max()[1], 2.0)
145            self.assertTrue (vg.min()[1] in (2.0, 0.0))
146        finally:
147            self.da.restoreGlobalVec(vg)
148            self.da.restoreLocalVec(vl)
149
150    def testGetOther(self):
151        ao = self.da.getAO()
152        lgmap = self.da.getLGMap()
153        l2g, g2l = self.da.getScatter()
154
155    def testRefineCoarsen(self):
156        da = self.da
157        rda = da.refine()
158        self.assertEqual(da.getDim(), rda.getDim())
159        self.assertEqual(da.getDof(), rda.getDof())
160        if da.dim != 1:
161            self.assertEqual(da.getStencilType(),  rda.getStencilType())
162        self.assertEqual(da.getStencilWidth(), rda.getStencilWidth())
163        cda = rda.coarsen()
164        self.assertEqual(rda.getDim(), cda.getDim())
165        self.assertEqual(rda.getDof(), cda.getDof())
166        for n1, n2 in zip(self.da.getSizes(), cda.getSizes()):
167            self.assertTrue(abs(n1-n2)<=1)
168
169    def testCoarsenRefine(self):
170        da = self.da
171        cda = self.da.coarsen()
172        self.assertEqual(da.getDim(), cda.getDim())
173        self.assertEqual(da.getDof(), cda.getDof())
174        if da.dim != 1:
175            self.assertEqual(da.getStencilType(),  cda.getStencilType())
176        self.assertEqual(da.getStencilWidth(), cda.getStencilWidth())
177        rda = cda.refine()
178        for n1, n2 in zip(self.da.getSizes(), rda.getSizes()):
179            self.assertTrue(abs(n1-n2)<=1)
180
181    def testRefineHierarchy(self):
182        levels = self.da.refineHierarchy(2)
183        self.assertTrue(isinstance(levels, list))
184        self.assertEqual(len(levels), 2)
185        for item in levels:
186            self.assertTrue(isinstance(item, PETSc.DM))
187
188    def testCoarsenHierarchy(self):
189        levels = self.da.coarsenHierarchy(2)
190        self.assertTrue(isinstance(levels, list))
191        self.assertEqual(len(levels), 2)
192        for item in levels:
193            self.assertTrue(isinstance(item, PETSc.DM))
194
195    def testCreateInterpolation(self):
196        da = self.da
197        if da.dim == 1: return
198        rda = da.refine()
199        mat, vec = da.createInterpolation(rda)
200
201    def testCreateInjection(self):
202        da = self.da
203        if da.dim == 1: return
204        rda = da.refine()
205        scatter = da.createInjection(rda)
206
207
208MIRROR   = PETSc.DMDA.BoundaryType.MIRROR
209GHOSTED  = PETSc.DMDA.BoundaryType.GHOSTED
210PERIODIC = PETSc.DMDA.BoundaryType.PERIODIC
211TWIST    = PETSc.DMDA.BoundaryType.TWIST
212
213SCALE = 4
214
215class BaseTestDA_1D(BaseTestDA):
216    SIZES = [100*SCALE]
217
218class BaseTestDA_2D(BaseTestDA):
219    SIZES = [9*SCALE,11*SCALE]
220
221class BaseTestDA_3D(BaseTestDA):
222    SIZES = [6*SCALE,7*SCALE,8*SCALE]
223
224# --------------------------------------------------------------------
225
226class TestDA_1D(BaseTestDA_1D, unittest.TestCase):
227    pass
228class TestDA_1D_W0(TestDA_1D):
229    SWIDTH = 0
230class TestDA_1D_W2(TestDA_1D):
231    SWIDTH = 2
232
233class TestDA_2D(BaseTestDA_2D, unittest.TestCase):
234    pass
235class TestDA_2D_W0(TestDA_2D):
236    SWIDTH = 0
237class TestDA_2D_W0_N2(TestDA_2D):
238    DOF = 2
239    SWIDTH = 0
240class TestDA_2D_W2(TestDA_2D):
241    SWIDTH = 2
242class TestDA_2D_W2_N2(TestDA_2D):
243    DOF = 2
244    SWIDTH = 2
245class TestDA_2D_PXY(TestDA_2D):
246    SIZES = [13*SCALE,17*SCALE]
247    DOF = 2
248    SWIDTH = 5
249    BOUNDARY = (PERIODIC,)*2
250class TestDA_2D_GXY(TestDA_2D):
251    SIZES = [13*SCALE,17*SCALE]
252    DOF = 2
253    SWIDTH = 5
254    BOUNDARY = (GHOSTED,)*2
255class TestDA_2D_TXY(TestDA_2D):
256    SIZES = [13*SCALE,17*SCALE]
257    DOF = 2
258    SWIDTH = 5
259    BOUNDARY = (TWIST,)*2
260
261class TestDA_3D(BaseTestDA_3D, unittest.TestCase):
262    pass
263class TestDA_3D_W0(TestDA_3D):
264    SWIDTH = 0
265class TestDA_3D_W0_N2(TestDA_3D):
266    DOF = 2
267    SWIDTH = 0
268class TestDA_3D_W2(TestDA_3D):
269    SWIDTH = 2
270class TestDA_3D_W2_N2(TestDA_3D):
271    DOF = 2
272    SWIDTH = 2
273class TestDA_3D_PXYZ(TestDA_3D):
274    SIZES = [11*SCALE,13*SCALE,17*SCALE]
275    DOF = 2
276    SWIDTH = 3
277    BOUNDARY = (PERIODIC,)*3
278class TestDA_3D_GXYZ(TestDA_3D):
279    SIZES = [11*SCALE,13*SCALE,17*SCALE]
280    DOF = 2
281    SWIDTH = 3
282    BOUNDARY = (GHOSTED,)*3
283class TestDA_3D_TXYZ(TestDA_3D):
284    SIZES = [11*SCALE,13*SCALE,17*SCALE]
285    DOF = 2
286    SWIDTH = 3
287    BOUNDARY = (TWIST,)*3
288
289# --------------------------------------------------------------------
290
291DIM = (1,2,3,)
292DOF = (None,1,2,3,4,5,)
293BOUNDARY_TYPE = (
294    None,
295    "none",     (0,)*3,        0,
296    "ghosted",  (GHOSTED,)*3,  GHOSTED,
297    "periodic", (PERIODIC,)*3, PERIODIC,
298    "twist",    (TWIST,)*3,    TWIST,
299    )
300STENCIL_TYPE  = (None,"star","box")
301STENCIL_WIDTH = (None,0,1,2,3)
302
303
304DIM           = (1,2,3)
305DOF           = (None,2,5)
306BOUNDARY_TYPE = (None,"none","periodic","ghosted","twist")
307STENCIL_TYPE  = (None,"box")
308STENCIL_WIDTH = (None,1,2)
309
310class TestDACreate(unittest.TestCase):
311    pass
312counter = 0
313for dim in DIM:
314    for dof in DOF:
315        for boundary in BOUNDARY_TYPE:
316            if isinstance(boundary, tuple):
317                boundary = boundary[:dim]
318            for stencil in STENCIL_TYPE:
319                for width in STENCIL_WIDTH:
320                    kargs = dict(sizes=[8*SCALE]*dim,
321                                 dim=dim, dof=dof,
322                                 boundary_type=boundary,
323                                 stencil_type=stencil,
324                                 stencil_width=width)
325                    def testCreate(self, kargs=kargs):
326                        kargs = dict(kargs)
327                        da = PETSc.DMDA().create(**kargs)
328                        da.destroy()
329                    setattr(TestDACreate,
330                            "testCreate%04d"%counter,
331                            testCreate)
332                    del testCreate, kargs
333                    counter += 1
334del counter, dim, dof, boundary, stencil, width
335
336class TestDADuplicate(unittest.TestCase):
337    pass
338counter = 0
339for dim in DIM:
340    for dof in DOF:
341        for boundary in BOUNDARY_TYPE:
342            if isinstance(boundary, tuple):
343                boundary = boundary[:dim]
344            for stencil in STENCIL_TYPE:
345                for width in STENCIL_WIDTH:
346                    kargs = dict(dim=dim, dof=dof,
347                                 boundary_type=boundary,
348                                 stencil_type=stencil,
349                                 stencil_width=width)
350                    def testDuplicate(self, kargs=kargs):
351                        kargs = dict(kargs)
352                        dim = kargs.pop('dim')
353                        dof = kargs['dof']
354                        boundary = kargs['boundary_type']
355                        stencil = kargs['stencil_type']
356                        width = kargs['stencil_width']
357                        da = PETSc.DMDA().create([8*SCALE]*dim)
358                        newda = da.duplicate(**kargs)
359                        self.assertEqual(newda.dim, da.dim)
360                        self.assertEqual(newda.sizes, da.sizes)
361                        self.assertEqual(newda.proc_sizes,
362                                         da.proc_sizes)
363                        self.assertEqual(newda.ranges, da.ranges)
364                        self.assertEqual(newda.corners, da.corners)
365                        if (newda.boundary_type == da.boundary_type
366                            and
367                            newda.stencil_width == da.stencil_width):
368                            self.assertEqual(newda.ghost_ranges,
369                                             da.ghost_ranges)
370                            self.assertEqual(newda.ghost_corners,
371                                             da.ghost_corners)
372                        if dof is None:
373                            dof = da.dof
374                        if boundary is None:
375                            boundary = da.boundary_type
376                        elif boundary == "none":
377                            boundary = (0,) * dim
378                        elif boundary == "mirror":
379                            boundary = (MIRROR,) * dim
380                        elif boundary == "ghosted":
381                            boundary = (GHOSTED,) * dim
382                        elif boundary == "periodic":
383                            boundary = (PERIODIC,) * dim
384                        elif boundary == "twist":
385                            boundary = (TWIST,) * dim
386                        elif isinstance(boundary, int):
387                            boundary = (boundary,) * dim
388                        if stencil is None:
389                            stencil = da.stencil[0]
390                        if width is None:
391                            width = da.stencil_width
392                        self.assertEqual(newda.dof, dof)
393                        self.assertEqual(newda.boundary_type,
394                                         boundary)
395                        if dim == 1:
396                            self.assertEqual(newda.stencil,
397                                             (stencil, width))
398                        newda.destroy()
399                        da.destroy()
400                    setattr(TestDADuplicate,
401                            "testDuplicate%04d"%counter,
402                            testDuplicate)
403                    del testDuplicate, kargs
404                    counter += 1
405del counter, dim, dof, boundary, stencil, width
406
407# --------------------------------------------------------------------
408
409if PETSc.COMM_WORLD.getSize() > 1:
410    del TestDA_1D_W0
411    del TestDA_2D_W0, TestDA_2D_W0_N2
412    del TestDA_3D_W0, TestDA_3D_W0_N2
413
414# --------------------------------------------------------------------
415
416if __name__ == '__main__':
417    unittest.main()
418
419# --------------------------------------------------------------------
420