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