1*55a74a43SLisandro Dalcintry: range = xrange 2*55a74a43SLisandro Dalcinexcept: pass 3*55a74a43SLisandro Dalcin 4*55a74a43SLisandro Dalcinfrom petsc4py import PETSc 5*55a74a43SLisandro Dalcin 6*55a74a43SLisandro Dalcin# grid size and spacing 7*55a74a43SLisandro Dalcinm, n = 32, 32 8*55a74a43SLisandro Dalcinhx = 1.0/(m-1) 9*55a74a43SLisandro Dalcinhy = 1.0/(n-1) 10*55a74a43SLisandro Dalcin 11*55a74a43SLisandro Dalcin# create sparse matrix 12*55a74a43SLisandro DalcinA = PETSc.Mat() 13*55a74a43SLisandro DalcinA.create(PETSc.COMM_WORLD) 14*55a74a43SLisandro DalcinA.setSizes([m*n, m*n]) 15*55a74a43SLisandro DalcinA.setType('aij') # sparse 16*55a74a43SLisandro DalcinA.setPreallocationNNZ(5) 17*55a74a43SLisandro Dalcin 18*55a74a43SLisandro Dalcin# precompute values for setting 19*55a74a43SLisandro Dalcin# diagonal and non-diagonal entries 20*55a74a43SLisandro Dalcindiagv = 2.0/hx**2 + 2.0/hy**2 21*55a74a43SLisandro Dalcinoffdx = -1.0/hx**2 22*55a74a43SLisandro Dalcinoffdy = -1.0/hy**2 23*55a74a43SLisandro Dalcin 24*55a74a43SLisandro Dalcin# loop over owned block of rows on this 25*55a74a43SLisandro Dalcin# processor and insert entry values 26*55a74a43SLisandro DalcinIstart, Iend = A.getOwnershipRange() 27*55a74a43SLisandro Dalcinfor I in range(Istart, Iend) : 28*55a74a43SLisandro Dalcin A[I,I] = diagv 29*55a74a43SLisandro Dalcin i = I//n # map row number to 30*55a74a43SLisandro Dalcin j = I - i*n # grid coordinates 31*55a74a43SLisandro Dalcin if i> 0 : J = I-n; A[I,J] = offdx 32*55a74a43SLisandro Dalcin if i< m-1: J = I+n; A[I,J] = offdx 33*55a74a43SLisandro Dalcin if j> 0 : J = I-1; A[I,J] = offdy 34*55a74a43SLisandro Dalcin if j< n-1: J = I+1; A[I,J] = offdy 35*55a74a43SLisandro Dalcin 36*55a74a43SLisandro Dalcin# communicate off-processor values 37*55a74a43SLisandro Dalcin# and setup internal data structures 38*55a74a43SLisandro Dalcin# for performing parallel operations 39*55a74a43SLisandro DalcinA.assemblyBegin() 40*55a74a43SLisandro DalcinA.assemblyEnd() 41