1# ------------------------------------------------------------------------ 2# 3# Poisson problem. This problem is modeled by the partial 4# differential equation 5# 6# -Laplacian(u) = 1, 0 < x,y < 1, 7# 8# with boundary conditions 9# 10# u = 0 for x = 0, x = 1, y = 0, y = 1 11# 12# A finite difference approximation with the usual 5-point stencil 13# is used to discretize the boundary value problem to obtain a 14# nonlinear system of equations. The problem is solved in a 2D 15# rectangular domain, using distributed arrays (DAs) to partition 16# the parallel grid. 17# 18# ------------------------------------------------------------------------ 19 20try: range = xrange 21except: pass 22 23import sys, petsc4py 24petsc4py.init(sys.argv) 25 26from petsc4py import PETSc 27 28class Poisson2D(object): 29 30 def __init__(self, da): 31 assert da.getDim() == 2 32 self.da = da 33 self.localX = da.createLocalVec() 34 35 def formRHS(self, B): 36 b = self.da.getVecArray(B) 37 mx, my = self.da.getSizes() 38 hx, hy = [1.0/m for m in [mx, my]] 39 (xs, xe), (ys, ye) = self.da.getRanges() 40 for j in range(ys, ye): 41 for i in range(xs, xe): 42 b[i, j] = 1*hx*hy 43 44 def mult(self, mat, X, Y): 45 # 46 self.da.globalToLocal(X, self.localX) 47 x = self.da.getVecArray(self.localX) 48 y = self.da.getVecArray(Y) 49 # 50 mx, my = self.da.getSizes() 51 hx, hy = [1.0/m for m in [mx, my]] 52 (xs, xe), (ys, ye) = self.da.getRanges() 53 for j in range(ys, ye): 54 for i in range(xs, xe): 55 u = x[i, j] # center 56 u_e = u_w = u_n = u_s = 0 57 if i > 0: u_w = x[i-1, j] # west 58 if i < mx-1: u_e = x[i+1, j] # east 59 if j > 0: u_s = x[i, j-1] # south 60 if j < ny-1: u_n = x[i, j+1] # north 61 u_xx = (-u_e + 2*u - u_w)*hy/hx 62 u_yy = (-u_n + 2*u - u_s)*hx/hy 63 y[i, j] = u_xx + u_yy 64 65OptDB = PETSc.Options() 66 67n = OptDB.getInt('n', 16) 68nx = OptDB.getInt('nx', n) 69ny = OptDB.getInt('ny', n) 70 71da = PETSc.DMDA().create([nx, ny], stencil_width=1, setup=False) 72da.setFromOptions() 73da.setUp() 74pde = Poisson2D(da) 75 76x = da.createGlobalVec() 77b = da.createGlobalVec() 78# A = da.createMat('python') 79A = PETSc.Mat().createPython( 80 [x.getSizes(), b.getSizes()], comm=da.comm) 81A.setPythonContext(pde) 82A.setUp() 83 84ksp = PETSc.KSP().create() 85ksp.setOperators(A) 86ksp.setType('cg') 87pc = ksp.getPC() 88pc.setType('none') 89ksp.setFromOptions() 90 91pde.formRHS(b) 92ksp.solve(b, x) 93 94u = da.createNaturalVec() 95da.globalToNatural(x, u) 96 97if OptDB.getBool('plot', True): 98 draw = PETSc.Viewer.DRAW(x.comm) 99 OptDB['draw_pause'] = 1 100 draw(x) 101