1import sys, petsc4py 2petsc4py.init(sys.argv) 3from petsc4py import PETSc 4 5# this user class is an application 6# context for the nonlinear problem 7# at hand; it contains some parameters 8# and knows how to compute residuals 9 10class Bratu2D: 11 12 def __init__(self, nx, ny, alpha, impl='python'): 13 self.nx = nx # x grid size 14 self.ny = ny # y grid size 15 self.alpha = alpha 16 if impl == 'python': 17 from bratu2dnpy import bratu2d 18 order = 'c' 19 elif impl == 'fortran': 20 from bratu2df90 import bratu2d 21 order = 'f' 22 else: 23 raise ValueError('invalid implementation') 24 self.compute = bratu2d 25 self.order = order 26 27 def evalFunction(self, snes, X, F): 28 nx, ny = self.nx, self.ny 29 alpha = self.alpha 30 order = self.order 31 x = X.getArray(readonly=1).reshape(nx, ny, order=order) 32 f = F.getArray(readonly=0).reshape(nx, ny, order=order) 33 self.compute(alpha, x, f) 34 35# convenience access to 36# PETSc options database 37OptDB = PETSc.Options() 38 39nx = OptDB.getInt('nx', 32) 40ny = OptDB.getInt('ny', nx) 41alpha = OptDB.getReal('alpha', 6.8) 42impl = OptDB.getString('impl', 'python') 43 44# create application context 45# and PETSc nonlinear solver 46appc = Bratu2D(nx, ny, alpha, impl) 47snes = PETSc.SNES().create() 48 49# register the function in charge of 50# computing the nonlinear residual 51f = PETSc.Vec().createSeq(nx*ny) 52snes.setFunction(appc.evalFunction, f) 53 54# configure the nonlinear solver 55# to use a matrix-free Jacobian 56snes.setUseMF(True) 57snes.getKSP().setType('cg') 58snes.setFromOptions() 59 60# solve the nonlinear problem 61b, x = None, f.duplicate() 62x.set(0) # zero initial guess 63snes.solve(b, x) 64 65if OptDB.getBool('plot', True): 66 da = PETSc.DMDA().create([nx,ny]) 67 u = da.createGlobalVec() 68 x.copy(u) 69 draw = PETSc.Viewer.DRAW() 70 OptDB['draw_pause'] = 1 71 draw(u) 72 73if OptDB.getBool('plot_mpl', False): 74 try: 75 from matplotlib import pylab 76 except ImportError: 77 PETSc.Sys.Print("matplotlib not available") 78 else: 79 from numpy import mgrid 80 X, Y = mgrid[0:1:1j*nx,0:1:1j*ny] 81 Z = x[...].reshape(nx,ny) 82 pylab.figure() 83 pylab.contourf(X,Y,Z) 84 pylab.colorbar() 85 pylab.plot(X.ravel(),Y.ravel(),'.k') 86 pylab.axis('equal') 87 pylab.show() 88