1import sys, petsc4py 2petsc4py.init(sys.argv) 3 4from petsc4py import PETSc 5from del2mat import Del2Mat 6 7# this a sequential example 8assert PETSc.COMM_WORLD.getSize() == 1 9 10# number of nodes in each direction 11# excluding those at the boundary 12n = 32 13h = 1.0/(n+1) # grid spacing 14 15# setup linear system matrix 16A = PETSc.Mat().create() 17A.setSizes([n**3, n**3]) 18A.setType('python') 19shell = Del2Mat(n) # shell context 20A.setPythonContext(shell) 21A.setUp() 22 23# setup linear system vectors 24x, b = A.createVecs() 25x.set(0.0) 26b.set(1.0) 27 28# setup Krylov solver 29ksp = PETSc.KSP().create() 30pc = ksp.getPC() 31ksp.setType('cg') 32pc.setType('none') 33 34# iteratively solve linear 35# system of equations A*x=b 36ksp.setOperators(A) 37ksp.setFromOptions() 38ksp.solve(b, x) 39 40# scale solution vector to 41# account for grid spacing 42x.scale(h**2) 43 44OptDB = PETSc.Options() 45if OptDB.getBool('plot_mpl', False): 46 try: 47 from matplotlib import pylab 48 except ImportError: 49 PETSc.Sys.Print("matplotlib not available") 50 else: 51 from numpy import mgrid 52 X, Y = mgrid[0:1:1j*n,0:1:1j*n] 53 Z = x[...].reshape(n,n,n)[:,:,n/2-2] 54 pylab.contourf(X, Y, Z) 55 pylab.axis('equal') 56 pylab.colorbar() 57 pylab.show() 58