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