1import sys, petsc4py 2petsc4py.init(sys.argv) 3 4from petsc4py import PETSc 5import Bratu3D as Bratu3D 6 7OptDB = PETSc.Options() 8 9N = OptDB.getInt('N', 16) 10lambda_ = OptDB.getReal('lambda', 6.0) 11do_plot = OptDB.getBool('plot', False) 12 13da = PETSc.DMDA().create([N, N, N], stencil_width=1) 14#app = App(da, lambda_) 15 16snes = PETSc.SNES().create() 17F = da.createGlobalVec() 18snes.setFunction(Bratu3D.formFunction, F, 19 args=(da, lambda_)) 20J = da.createMat() 21snes.setJacobian(Bratu3D.formJacobian, J, 22 args=(da, lambda_)) 23 24snes.setFromOptions() 25 26X = da.createGlobalVec() 27Bratu3D.formInitGuess(X, da, lambda_) 28snes.solve(None, X) 29 30U = da.createNaturalVec() 31da.globalToNatural(X, U) 32 33 34def plot(da, U): 35 scatter, U0 = PETSc.Scatter.toZero(U) 36 scatter.scatter(U, U0, False, PETSc.Scatter.Mode.FORWARD) 37 rank = PETSc.COMM_WORLD.getRank() 38 if rank == 0: 39 solution = U0[...].reshape(da.sizes, order='f').copy() 40 try: 41 from matplotlib import pyplot 42 pyplot.contourf(solution[:, :, N//2]) 43 pyplot.axis('equal') 44 pyplot.show() 45 except: 46 pass 47 PETSc.COMM_WORLD.barrier() 48 scatter.destroy() 49 U0.destroy() 50 51 52if do_plot: plot(da, U) 53 54 55U.destroy() 56X.destroy() 57F.destroy() 58J.destroy() 59da.destroy() 60snes.destroy() 61