xref: /petsc/src/binding/petsc4py/test/test_dmplex.py (revision 1d4815f0dddabe37e106113de1c05b2effb59878)
1import petsc4py
2from petsc4py import PETSc
3import unittest
4import os
5import filecmp
6import numpy as np
7
8# --------------------------------------------------------------------
9
10ERR_SUP = 56
11
12class BaseTestPlex(object):
13
14    COMM = PETSc.COMM_WORLD
15    DIM = 1
16    CELLS = [[0, 1], [1, 2]]
17    COORDS = [[0.], [0.5], [1.]]
18    COMP = 1
19    DOFS = [1, 0]
20
21    def setUp(self):
22        self.plex = PETSc.DMPlex().createFromCellList(self.DIM,
23                                                      self.CELLS,
24                                                      self.COORDS,
25                                                      comm=self.COMM)
26
27    def tearDown(self):
28        self.plex.destroy()
29        self.plex = None
30
31    def testTopology(self):
32        dim = self.plex.getDimension()
33        pStart, pEnd = self.plex.getChart()
34        cStart, cEnd = self.plex.getHeightStratum(0)
35        vStart, vEnd = self.plex.getDepthStratum(0)
36        numDepths = self.plex.getLabelSize("depth")
37        coords_raw = self.plex.getCoordinates().getArray()
38        coords = np.reshape(coords_raw, (vEnd - vStart, dim))
39        self.assertEqual(dim, self.DIM)
40        self.assertEqual(numDepths, self.DIM+1)
41        if self.CELLS is not None:
42            self.assertEqual(cEnd-cStart, len(self.CELLS))
43        if self.COORDS is not None:
44            self.assertEqual(vEnd-vStart, len(self.COORDS))
45            self.assertTrue((coords == self.COORDS).all())
46
47    def testClosure(self):
48        pStart, pEnd = self.plex.getChart()
49        for p in range(pStart, pEnd):
50            closure = self.plex.getTransitiveClosure(p)[0]
51            for c in closure:
52                cone = self.plex.getCone(c)
53                self.assertEqual(self.plex.getConeSize(c), len(cone))
54                for i in cone:
55                    self.assertIn(i, closure)
56            star = self.plex.getTransitiveClosure(p, useCone=False)[0]
57            for s in star:
58                support = self.plex.getSupport(s)
59                self.assertEqual(self.plex.getSupportSize(s), len(support))
60                for i in support:
61                    self.assertIn(i, star)
62
63    def testAdjacency(self):
64        PETSc.DMPlex.setAdjacencyUseAnchors(self.plex, False)
65        flag = PETSc.DMPlex.getAdjacencyUseAnchors(self.plex)
66        self.assertFalse(flag)
67        PETSc.DMPlex.setAdjacencyUseAnchors(self.plex, True)
68        flag = PETSc.DMPlex.getAdjacencyUseAnchors(self.plex)
69        self.assertTrue(flag)
70        PETSc.DMPlex.setBasicAdjacency(self.plex, False, False)
71        flagA, flagB = PETSc.DMPlex.getBasicAdjacency(self.plex)
72        self.assertFalse(flagA)
73        self.assertFalse(flagB)
74        PETSc.DMPlex.setBasicAdjacency(self.plex, True, True)
75        flagA, flagB = PETSc.DMPlex.getBasicAdjacency(self.plex)
76        self.assertTrue(flagA)
77        self.assertTrue(flagB)
78        pStart, pEnd = self.plex.getChart()
79        for p in range(pStart, pEnd):
80            adjacency = self.plex.getAdjacency(p)
81            self.assertTrue(p in adjacency)
82            self.assertTrue(len(adjacency) > 1)
83
84    def testSectionDofs(self):
85        self.plex.setNumFields(1)
86        section = self.plex.createSection([self.COMP], [self.DOFS])
87        size = section.getStorageSize()
88        entity_dofs = [self.plex.getStratumSize("depth", d) *
89                       self.DOFS[d] for d in range(self.DIM+1)]
90        self.assertEqual(sum(entity_dofs), size)
91
92    def testSectionClosure(self):
93        section = self.plex.createSection([self.COMP], [self.DOFS])
94        self.plex.setSection(section)
95        vec = self.plex.createLocalVec()
96        pStart, pEnd = self.plex.getChart()
97        for p in range(pStart, pEnd):
98            for i in range(section.getDof(p)):
99                off = section.getOffset(p)
100                vec.setValue(off+i, p)
101
102        for p in range(pStart, pEnd):
103            point_closure = self.plex.getTransitiveClosure(p)[0]
104            dof_closure = self.plex.vecGetClosure(section, vec, p)
105            for p in dof_closure:
106                self.assertIn(p, point_closure)
107
108    def testBoundaryLabel(self):
109        pStart, pEnd = self.plex.getChart()
110        if (pEnd - pStart == 0): return
111
112        self.assertFalse(self.plex.hasLabel("boundary"))
113        self.plex.markBoundaryFaces("boundary")
114        self.assertTrue(self.plex.hasLabel("boundary"))
115
116        faces = self.plex.getStratumIS("boundary", 1)
117        for f in faces.getIndices():
118            points, orient = self.plex.getTransitiveClosure(f, useCone=True)
119            for p in points:
120                self.plex.setLabelValue("boundary", p, 1)
121
122        for p in range(pStart, pEnd):
123            if self.plex.getLabelValue("boundary", p) != 1:
124                self.plex.setLabelValue("boundary", p, 2)
125
126        numBoundary = self.plex.getStratumSize("boundary", 1)
127        numInterior = self.plex.getStratumSize("boundary", 2)
128        self.assertNotEqual(numBoundary, pEnd - pStart)
129        self.assertNotEqual(numInterior, pEnd - pStart)
130        self.assertEqual(numBoundary + numInterior, pEnd - pStart)
131
132
133    def testAdapt(self):
134        dim = self.plex.getDimension()
135        if dim == 1: return
136        vStart, vEnd = self.plex.getDepthStratum(0)
137        numVertices = vEnd-vStart
138        metric_array = np.zeros([numVertices,dim,dim])
139        for met in metric_array:
140            met[:,:] = np.diag([9]*dim)
141        metric = PETSc.Vec().createWithArray(metric_array)
142        try:
143            newplex = self.plex.adaptMetric(metric,"")
144        except PETSc.Error as exc:
145            if exc.ierr != ERR_SUP: raise
146
147
148# --------------------------------------------------------------------
149
150class BaseTestPlex_2D(BaseTestPlex):
151    DIM = 2
152    CELLS = [[0, 1, 3], [1, 3, 4], [1, 2, 4], [2, 4, 5],
153             [3, 4, 6], [4, 6, 7], [4, 5, 7], [5, 7, 8]]
154    COORDS = [[0.0, 0.0], [0.5, 0.0], [1.0, 0.0],
155              [0.0, 0.5], [0.5, 0.5], [1.0, 0.5],
156              [0.0, 1.0], [0.5, 1.0], [1.0, 1.0]]
157    DOFS = [1, 0, 0]
158
159class BaseTestPlex_3D(BaseTestPlex):
160    DIM = 3
161    CELLS = [[0, 2, 3, 7], [0, 2, 6, 7], [0, 4, 6, 7],
162             [0, 1, 3, 7], [0, 1, 5, 7], [0, 4, 5, 7]]
163    COORDS = [[0., 0., 0.], [1., 0., 0.], [0., 1., 0.], [1., 1., 0.],
164              [0., 0., 1.], [1., 0., 1.], [0., 1., 1.], [1., 1., 1.]]
165    DOFS = [1, 0, 0, 0]
166
167# --------------------------------------------------------------------
168
169class TestPlex_1D(BaseTestPlex, unittest.TestCase):
170    pass
171
172class TestPlex_2D(BaseTestPlex_2D, unittest.TestCase):
173    pass
174
175class TestPlex_3D(BaseTestPlex_3D, unittest.TestCase):
176    pass
177
178class TestPlex_2D_P3(BaseTestPlex_2D, unittest.TestCase):
179    DOFS = [1, 2, 1]
180
181class TestPlex_3D_P3(BaseTestPlex_3D, unittest.TestCase):
182    DOFS = [1, 2, 1, 0]
183
184class TestPlex_3D_P4(BaseTestPlex_3D, unittest.TestCase):
185    DOFS = [1, 3, 3, 1]
186
187class TestPlex_2D_BoxTensor(BaseTestPlex_2D, unittest.TestCase):
188    CELLS = None
189    COORDS = None
190    def setUp(self):
191        self.plex = PETSc.DMPlex().createBoxMesh([3,3], simplex=False)
192
193class TestPlex_3D_BoxTensor(BaseTestPlex_3D, unittest.TestCase):
194    CELLS = None
195    COORDS = None
196    def setUp(self):
197        self.plex = PETSc.DMPlex().createBoxMesh([3,3,3], simplex=False)
198
199try:
200    raise PETSc.Error
201    PETSc.DMPlex().createBoxMesh([2,2], simplex=True, comm=PETSc.COMM_SELF).destroy()
202except PETSc.Error:
203    pass
204else:
205    class TestPlex_2D_Box(BaseTestPlex_2D, unittest.TestCase):
206        CELLS = None
207        COORDS = None
208        def setUp(self):
209            self.plex = PETSc.DMPlex().createBoxMesh([1,1], simplex=True)
210
211    class TestPlex_2D_Boundary(BaseTestPlex_2D, unittest.TestCase):
212        CELLS = None
213        COORDS = None
214        def setUp(self):
215            boundary = PETSc.DMPlex().create(self.COMM)
216            boundary.createSquareBoundary([0., 0.], [1., 1.], [2, 2])
217            boundary.setDimension(self.DIM-1)
218            self.plex = PETSc.DMPlex().generate(boundary)
219
220    class TestPlex_3D_Box(BaseTestPlex_3D, unittest.TestCase):
221        CELLS = None
222        COORDS = None
223        def setUp(self):
224            self.plex = PETSc.DMPlex().createBoxMesh([1,1,1], simplex=True)
225
226    class TestPlex_3D_Boundary(BaseTestPlex_3D, unittest.TestCase):
227        CELLS = None
228        COORDS = None
229        def setUp(self):
230            boundary = PETSc.DMPlex().create(self.COMM)
231            boundary.createCubeBoundary([0., 0., 0.], [1., 1., 1.], [1, 1, 1])
232            boundary.setDimension(self.DIM-1)
233            self.plex = PETSc.DMPlex().generate(boundary)
234
235# --------------------------------------------------------------------
236
237PETSC_DIR = petsc4py.get_config()['PETSC_DIR']
238
239def check_dtype(method):
240    def wrapper(self, *args, **kwargs):
241        if PETSc.ScalarType is PETSc.ComplexType:
242            return
243        else:
244            return method(self, *args, **kwargs)
245    return wrapper
246
247def check_package(method):
248    def wrapper(self, *args, **kwargs):
249        if not PETSc.Sys.hasExternalPackage("hdf5"):
250            return
251        elif self.PARTITIONERTYPE != "simple" and \
252           not PETSc.Sys.hasExternalPackage(self.PARTITIONERTYPE):
253            return
254        else:
255            return method(self, *args, **kwargs)
256    return wrapper
257
258def check_nsize(method):
259    def wrapper(self, *args, **kwargs):
260        if PETSc.COMM_WORLD.size != self.NSIZE:
261            return
262        else:
263            return method(self, *args, **kwargs)
264    return wrapper
265
266class BaseTestPlexHDF5(object):
267    NSIZE = 4
268    NTIMES = 3
269
270    def setUp(self):
271        self.txtvwr = PETSc.Viewer()
272
273    def tearDown(self):
274        if not PETSc.COMM_WORLD.rank:
275            if os.path.exists(self.outfile()):
276                os.remove(self.outfile())
277            if os.path.exists(self.tmp_output_file()):
278                os.remove(self.tmp_output_file())
279        self.txtvwr = None
280
281    def _name(self):
282        return "%s_outformat-%s_%s" % (self.SUFFIX,
283                                       self.OUTFORMAT,
284                                       self.PARTITIONERTYPE)
285
286    def infile(self):
287        return os.path.join(PETSC_DIR, "share/petsc/datafiles/",
288                            "meshes/blockcylinder-50.h5")
289
290    def outfile(self):
291        return os.path.join("./temp_test_dmplex_%s.h5" % self._name())
292
293    def informat(self):
294        return PETSc.Viewer.Format.HDF5_XDMF
295
296    def outformat(self):
297        d = {"hdf5_petsc": PETSc.Viewer.Format.HDF5_PETSC,
298             "hdf5_xdmf": PETSc.Viewer.Format.HDF5_XDMF}
299        return d[self.OUTFORMAT]
300
301    def partitionerType(self):
302        d = {"simple": PETSc.Partitioner.Type.SIMPLE,
303             "ptscotch": PETSc.Partitioner.Type.PTSCOTCH,
304             "parmetis": PETSc.Partitioner.Type.PARMETIS}
305        return d[self.PARTITIONERTYPE]
306
307    def ref_output_file(self):
308        return os.path.join(PETSC_DIR, "src/dm/impls/plex/tutorials/",
309                            "output/ex5_%s.out" % self._name())
310
311    def tmp_output_file(self):
312        return os.path.join("./temp_test_dmplex_%s.out" % self._name())
313
314    def outputText(self, msg, comm):
315        if not comm.rank:
316            with open(self.tmp_output_file(), 'a') as f:
317                f.write(msg)
318
319    def outputPlex(self, plex):
320        self.txtvwr.createASCII(self.tmp_output_file(),
321                                mode='a', comm=plex.comm)
322        plex.view(viewer=self.txtvwr)
323        self.txtvwr.destroy()
324
325    @check_dtype
326    @check_package
327    @check_nsize
328    def testViewLoadCycle(self):
329        grank = PETSc.COMM_WORLD.rank
330        for i in range(self.NTIMES):
331            if i == 0:
332                infname = self.infile()
333                informt = self.informat()
334            else:
335                infname = self.outfile()
336                informt = self.outformat()
337            if self.HETEROGENEOUS:
338                mycolor = (grank > self.NTIMES - i)
339            else:
340                mycolor = 0
341            mpicomm = PETSc.COMM_WORLD.tompi4py()
342            comm = PETSc.Comm(comm=mpicomm.Split(color=mycolor, key=grank))
343            if mycolor == 0:
344                self.outputText("Begin cycle %d\n" % i, comm)
345                plex = PETSc.DMPlex()
346                vwr = PETSc.ViewerHDF5()
347                # Create plex
348                plex.create(comm=comm)
349                plex.setName("DMPlex Object")
350                # Load data from XDMF into dm in parallel
351                vwr.create(infname, mode='r', comm=comm)
352                vwr.pushFormat(format=informt)
353                plex.load(viewer=vwr)
354                plex.setOptionsPrefix("loaded_")
355                plex.setFromOptions()
356                vwr.popFormat()
357                vwr.destroy()
358                self.outputPlex(plex)
359                # Test DM is indeed distributed
360                flg = plex.isDistributed()
361                self.outputText("Loaded mesh distributed? %s\n" %
362                                str(flg).upper(), comm)
363                # Interpolate
364                plex.interpolate()
365                plex.setOptionsPrefix("interpolated_")
366                plex.setFromOptions()
367                self.outputPlex(plex)
368                # Redistribute
369                part = plex.getPartitioner()
370                part.setType(self.partitionerType())
371                _ = plex.distribute(overlap=0)
372                plex.setOptionsPrefix("redistributed_")
373                plex.setFromOptions()
374                self.outputPlex(plex)
375                # Save redistributed dm to XDMF in parallel
376                vwr.create(self.outfile(), mode='w', comm=comm)
377                vwr.pushFormat(format=self.outformat())
378                plex.view(viewer=vwr)
379                vwr.popFormat()
380                vwr.destroy()
381                # Destroy plex
382                plex.destroy()
383                self.outputText("End   cycle %d\n--------\n" % i, comm)
384            PETSc.COMM_WORLD.Barrier()
385        # Check that the output is identical to that of plex/tutorial/ex5.c.
386        self.assertTrue(filecmp.cmp(self.tmp_output_file(),
387                                    self.ref_output_file(), shallow=False),
388                        'Contents of the files not the same.')
389        PETSc.COMM_WORLD.Barrier()
390
391class BaseTestPlexHDF5Homogeneous(BaseTestPlexHDF5):
392    """Test save on N / load on N."""
393    SUFFIX = 0
394    HETEROGENEOUS = False
395
396class BaseTestPlexHDF5Heterogeneous(BaseTestPlexHDF5):
397    """Test save on N / load on M."""
398    SUFFIX = 1
399    HETEROGENEOUS = True
400
401class TestPlexHDF5PETSCSimpleHomogeneous(BaseTestPlexHDF5Homogeneous,
402                                         unittest.TestCase):
403    OUTFORMAT = "hdf5_petsc"
404    PARTITIONERTYPE = "simple"
405
406"""
407Skipping. PTScotch produces different distributions when run
408in a sequence in a single session.
409
410class TestPlexHDF5PETSCPTScotchHomogeneous(BaseTestPlexHDF5Homogeneous,
411                                           unittest.TestCase):
412    OUTFORMAT = "hdf5_petsc"
413    PARTITIONERTYPE = "ptscotch"
414"""
415
416class TestPlexHDF5PETSCParmetisHomogeneous(BaseTestPlexHDF5Homogeneous,
417                                           unittest.TestCase):
418    OUTFORMAT = "hdf5_petsc"
419    PARTITIONERTYPE = "parmetis"
420
421class TestPlexHDF5XDMFSimpleHomogeneous(BaseTestPlexHDF5Homogeneous,
422                                        unittest.TestCase):
423    OUTFORMAT = "hdf5_xdmf"
424    PARTITIONERTYPE = "simple"
425
426"""
427Skipping. PTScotch produces different distributions when run
428in a sequence in a single session.
429
430class TestPlexHDF5XDMFPTScotchHomogeneous(BaseTestPlexHDF5Homogeneous,
431                                          unittest.TestCase):
432    OUTFORMAT = "hdf5_xdmf"
433    PARTITIONERTYPE = "ptscotch"
434"""
435
436class TestPlexHDF5XDMFParmetisHomogeneous(BaseTestPlexHDF5Homogeneous,
437                                          unittest.TestCase):
438    OUTFORMAT = "hdf5_xdmf"
439    PARTITIONERTYPE = "parmetis"
440
441class TestPlexHDF5PETSCSimpleHeterogeneous(BaseTestPlexHDF5Heterogeneous,
442                                           unittest.TestCase):
443    OUTFORMAT = "hdf5_petsc"
444    PARTITIONERTYPE = "simple"
445
446"""
447Skipping. PTScotch produces different distributions when run
448in a sequence in a single session.
449
450class TestPlexHDF5PETSCPTScotchHeterogeneous(BaseTestPlexHDF5Heterogeneous,
451                                             unittest.TestCase):
452    OUTFORMAT = "hdf5_petsc"
453    PARTITIONERTYPE = "ptscotch"
454"""
455
456class TestPlexHDF5PETSCParmetisHeterogeneous(BaseTestPlexHDF5Heterogeneous,
457                                             unittest.TestCase):
458    OUTFORMAT = "hdf5_petsc"
459    PARTITIONERTYPE = "parmetis"
460
461class TestPlexHDF5XDMFSimpleHeterogeneous(BaseTestPlexHDF5Heterogeneous,
462                                          unittest.TestCase):
463    OUTFORMAT = "hdf5_xdmf"
464    PARTITIONERTYPE = "simple"
465
466class TestPlexHDF5XDMFPTScotchHeterogeneous(BaseTestPlexHDF5Heterogeneous,
467                                            unittest.TestCase):
468    OUTFORMAT = "hdf5_xdmf"
469    PARTITIONERTYPE = "ptscotch"
470
471class TestPlexHDF5XDMFParmetisHeterogeneous(BaseTestPlexHDF5Heterogeneous,
472                                            unittest.TestCase):
473    OUTFORMAT = "hdf5_xdmf"
474    PARTITIONERTYPE = "parmetis"
475
476# --------------------------------------------------------------------
477
478if __name__ == '__main__':
479    unittest.main()
480