1from petsc4py import PETSc
2import unittest
3import numpy as np
4
5# --------------------------------------------------------------------
6
7ERR_SUP = 56
8
9class BaseTestPlex(object):
10
11    COMM = PETSc.COMM_WORLD
12    DIM = 1
13    CELLS = [[0, 1], [1, 2]]
14    COORDS = [[0.], [0.5], [1.]]
15    COMP = 1
16    DOFS = [1, 0]
17
18    def setUp(self):
19        self.plex = PETSc.DMPlex().createFromCellList(self.DIM,
20                                                      self.CELLS,
21                                                      self.COORDS,
22                                                      comm=self.COMM)
23
24    def tearDown(self):
25        self.plex.destroy()
26        self.plex = None
27
28    def testTopology(self):
29        dim = self.plex.getDimension()
30        pStart, pEnd = self.plex.getChart()
31        cStart, cEnd = self.plex.getHeightStratum(0)
32        vStart, vEnd = self.plex.getDepthStratum(0)
33        numDepths = self.plex.getLabelSize("depth")
34        coords_raw = self.plex.getCoordinates().getArray()
35        coords = np.reshape(coords_raw, (vEnd - vStart, dim))
36        self.assertEqual(dim, self.DIM)
37        self.assertEqual(numDepths, self.DIM+1)
38        if self.CELLS is not None:
39            self.assertEqual(cEnd-cStart, len(self.CELLS))
40        if self.COORDS is not None:
41            self.assertEqual(vEnd-vStart, len(self.COORDS))
42            self.assertTrue((coords == self.COORDS).all())
43
44    def testClosure(self):
45        pStart, pEnd = self.plex.getChart()
46        for p in range(pStart, pEnd):
47            closure = self.plex.getTransitiveClosure(p)[0]
48            for c in closure:
49                cone = self.plex.getCone(c)
50                self.assertEqual(self.plex.getConeSize(c), len(cone))
51                for i in cone:
52                    self.assertIn(i, closure)
53            star = self.plex.getTransitiveClosure(p, useCone=False)[0]
54            for s in star:
55                support = self.plex.getSupport(s)
56                self.assertEqual(self.plex.getSupportSize(s), len(support))
57                for i in support:
58                    self.assertIn(i, star)
59
60    def testAdjacency(self):
61        PETSc.DMPlex.setAdjacencyUseAnchors(self.plex, False)
62        flag = PETSc.DMPlex.getAdjacencyUseAnchors(self.plex)
63        self.assertFalse(flag)
64        PETSc.DMPlex.setAdjacencyUseAnchors(self.plex, True)
65        flag = PETSc.DMPlex.getAdjacencyUseAnchors(self.plex)
66        self.assertTrue(flag)
67        PETSc.DMPlex.setBasicAdjacency(self.plex, False, False)
68        flagA, flagB = PETSc.DMPlex.getBasicAdjacency(self.plex)
69        self.assertFalse(flagA)
70        self.assertFalse(flagB)
71        PETSc.DMPlex.setBasicAdjacency(self.plex, True, True)
72        flagA, flagB = PETSc.DMPlex.getBasicAdjacency(self.plex)
73        self.assertTrue(flagA)
74        self.assertTrue(flagB)
75        pStart, pEnd = self.plex.getChart()
76        for p in range(pStart, pEnd):
77            adjacency = self.plex.getAdjacency(p)
78            self.assertTrue(p in adjacency)
79            self.assertTrue(len(adjacency) > 1)
80
81    def testSectionDofs(self):
82        self.plex.setNumFields(1)
83        section = self.plex.createSection([self.COMP], [self.DOFS])
84        size = section.getStorageSize()
85        entity_dofs = [self.plex.getStratumSize("depth", d) *
86                       self.DOFS[d] for d in range(self.DIM+1)]
87        self.assertEqual(sum(entity_dofs), size)
88
89    def testSectionClosure(self):
90        section = self.plex.createSection([self.COMP], [self.DOFS])
91        self.plex.setSection(section)
92        vec = self.plex.createLocalVec()
93        pStart, pEnd = self.plex.getChart()
94        for p in range(pStart, pEnd):
95            for i in range(section.getDof(p)):
96                off = section.getOffset(p)
97                vec.setValue(off+i, p)
98
99        for p in range(pStart, pEnd):
100            point_closure = self.plex.getTransitiveClosure(p)[0]
101            dof_closure = self.plex.vecGetClosure(section, vec, p)
102            for p in dof_closure:
103                self.assertIn(p, point_closure)
104
105    def testBoundaryLabel(self):
106        self.assertFalse(self.plex.hasLabel("boundary"))
107        self.plex.markBoundaryFaces("boundary")
108        self.assertTrue(self.plex.hasLabel("boundary"))
109
110        faces = self.plex.getStratumIS("boundary", 1)
111        for f in faces.getIndices():
112            points, orient = self.plex.getTransitiveClosure(f, useCone=True)
113            for p in points:
114                self.plex.setLabelValue("boundary", p, 1)
115
116        pStart, pEnd = self.plex.getChart()
117        for p in range(pStart, pEnd):
118            if self.plex.getLabelValue("boundary", p) != 1:
119                self.plex.setLabelValue("boundary", p, 2)
120
121        numBoundary = self.plex.getStratumSize("boundary", 1)
122        numInterior = self.plex.getStratumSize("boundary", 2)
123        self.assertNotEqual(numBoundary, pEnd - pStart)
124        self.assertNotEqual(numInterior, pEnd - pStart)
125        self.assertEqual(numBoundary + numInterior, pEnd - pStart)
126
127
128    def testAdapt(self):
129        dim = self.plex.getDimension()
130        if dim == 1: return
131        vStart, vEnd = self.plex.getDepthStratum(0)
132        numVertices = vEnd-vStart
133        metric_array = np.zeros([numVertices,dim,dim])
134        for met in metric_array:
135            met[:,:] = np.diag([9]*dim)
136        metric = PETSc.Vec().createWithArray(metric_array)
137        try:
138            newplex = self.plex.adaptMetric(metric,"")
139        except PETSc.Error as exc:
140            if exc.ierr != ERR_SUP: raise
141
142
143# --------------------------------------------------------------------
144
145class BaseTestPlex_2D(BaseTestPlex):
146    DIM = 2
147    CELLS = [[0, 1, 3], [1, 3, 4], [1, 2, 4], [2, 4, 5],
148             [3, 4, 6], [4, 6, 7], [4, 5, 7], [5, 7, 8]]
149    COORDS = [[0.0, 0.0], [0.5, 0.0], [1.0, 0.0],
150              [0.0, 0.5], [0.5, 0.5], [1.0, 0.5],
151              [0.0, 1.0], [0.5, 1.0], [1.0, 1.0]]
152    DOFS = [1, 0, 0]
153
154class BaseTestPlex_3D(BaseTestPlex):
155    DIM = 3
156    CELLS = [[0, 2, 3, 7], [0, 2, 6, 7], [0, 4, 6, 7],
157             [0, 1, 3, 7], [0, 1, 5, 7], [0, 4, 5, 7]]
158    COORDS = [[0., 0., 0.], [1., 0., 0.], [0., 1., 0.], [1., 1., 0.],
159              [0., 0., 1.], [1., 0., 1.], [0., 1., 1.], [1., 1., 1.]]
160    DOFS = [1, 0, 0, 0]
161
162# --------------------------------------------------------------------
163
164class TestPlex_1D(BaseTestPlex, unittest.TestCase):
165    pass
166
167class TestPlex_2D(BaseTestPlex_2D, unittest.TestCase):
168    pass
169
170class TestPlex_3D(BaseTestPlex_3D, unittest.TestCase):
171    pass
172
173class TestPlex_2D_P3(BaseTestPlex_2D, unittest.TestCase):
174    DOFS = [1, 2, 1]
175
176class TestPlex_3D_P3(BaseTestPlex_3D, unittest.TestCase):
177    DOFS = [1, 2, 1, 0]
178
179class TestPlex_3D_P4(BaseTestPlex_3D, unittest.TestCase):
180    DOFS = [1, 3, 3, 1]
181
182class TestPlex_2D_BoxTensor(BaseTestPlex_2D, unittest.TestCase):
183    CELLS = None
184    COORDS = None
185    def setUp(self):
186        self.plex = PETSc.DMPlex().createBoxMesh([3,3], simplex=False)
187
188class TestPlex_3D_BoxTensor(BaseTestPlex_3D, unittest.TestCase):
189    CELLS = None
190    COORDS = None
191    def setUp(self):
192        self.plex = PETSc.DMPlex().createBoxMesh([3,3,3], simplex=False)
193
194import sys
195try:
196    raise PETSc.Error
197    PETSc.DMPlex().createBoxMesh([2,2], simplex=True, comm=PETSc.COMM_SELF).destroy()
198except PETSc.Error:
199    pass
200else:
201    class TestPlex_2D_Box(BaseTestPlex_2D, unittest.TestCase):
202        CELLS = None
203        COORDS = None
204        def setUp(self):
205            self.plex = PETSc.DMPlex().createBoxMesh([1,1], simplex=True)
206
207    class TestPlex_2D_Boundary(BaseTestPlex_2D, unittest.TestCase):
208        CELLS = None
209        COORDS = None
210        def setUp(self):
211            boundary = PETSc.DMPlex().create(self.COMM)
212            boundary.createSquareBoundary([0., 0.], [1., 1.], [2, 2])
213            boundary.setDimension(self.DIM-1)
214            self.plex = PETSc.DMPlex().generate(boundary)
215
216    class TestPlex_3D_Box(BaseTestPlex_3D, unittest.TestCase):
217        CELLS = None
218        COORDS = None
219        def setUp(self):
220            self.plex = PETSc.DMPlex().createBoxMesh([1,1,1], simplex=True)
221
222    class TestPlex_3D_Boundary(BaseTestPlex_3D, unittest.TestCase):
223        CELLS = None
224        COORDS = None
225        def setUp(self):
226            boundary = PETSc.DMPlex().create(self.COMM)
227            boundary.createCubeBoundary([0., 0., 0.], [1., 1., 1.], [1, 1, 1])
228            boundary.setDimension(self.DIM-1)
229            self.plex = PETSc.DMPlex().generate(boundary)
230
231# --------------------------------------------------------------------
232
233if __name__ == '__main__':
234    unittest.main()
235