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