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