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