1# -*- coding: utf-8 -*-
2"""
3Created on Fri Dec 25 17:03:18 2015
4
5@author: ale
6"""
7
8import sys,petsc4py
9petsc4py.init(sys.argv)
10from petsc4py import PETSc
11import numpy as np
12
13dim = 2
14if not PETSc.COMM_WORLD.rank:
15    coords = np.asarray([[0.0, 0.0],
16                         [0.5, 0.0],
17                         [1.0, 0.0],
18                         [0.0, 0.5],
19                         [0.5, 0.5],
20                         [1.0, 0.5],
21                         [0.0, 1.0],
22                         [0.5, 1.0],
23                         [1.0, 1.0]], dtype=float)
24    cells = np.asarray([[0,1,4,3],
25                        [1,2,5,4],
26                        [3,4,7,6],
27                        [4,5,8,7]], dtype=PETSc.IntType)
28else:
29    coords = np.zeros((0, 2), dtype=float)
30    cells = np.zeros((0, 4), dtype=PETSc.IntType)
31
32plex = PETSc.DMPlex().createFromCellList(dim, cells, coords, comm=PETSc.COMM_WORLD)
33
34pStart, pEnd = plex.getChart()
35plex.view()
36print("pStart, pEnd: ", pStart, pEnd)
37
38# Create section with 1 field with 1 DoF per vertex, edge amd cell
39numComp = 1
40# Start with an empty vector
41numDof = [0] * 3
42# Field defined on vertexes
43numDof[0] = 1
44# Field defined on edges
45numDof[1] = 1
46# Field defined on cells
47numDof[2] = 1
48
49plex.setNumFields(1)
50origSect = plex.createSection(numComp, numDof)
51origSect.setFieldName(0, 'TestField')
52origSect.setUp()
53origSect.view()
54
55plex.setSection(origSect)
56origVec = plex.createGlobalVec()
57origVec.view()
58
59origVec.setValues(list(range(pStart, pEnd)),list(range(pStart,pEnd)))
60origVec.assemblyBegin()
61origVec.assemblyEnd()
62
63origVec.view()
64
65if PETSc.COMM_WORLD.size > 1:
66    sf = plex.distribute()
67    sf.view()
68
69    newSect, newVec = plex.distributeField(sf, origSect, origVec)
70
71else:
72    newSect = origSect
73    newVec = origVec
74
75newSect.view()
76newVec.view()
77
78