1import sys, petsc4py 2petsc4py.init(sys.argv) 3 4from petsc4py import PETSc 5import Bratu3D as Bratu3D 6 7class App(object): 8 9 def __init__(self, da, lambda_): 10 assert da.getDim() == 3 11 self.da = da 12 self.params = Bratu3D.Params() 13 self.params.lambda_ = lambda_ 14 15 def formInitGuess(self, X): 16 X.zeroEntries() # just in case 17 Bratu3D.FormInitGuess(self.da, X, self.params) 18 19 def formFunction(self, snes, X, F): 20 F.zeroEntries() # just in case 21 Bratu3D.FormFunction(self.da, X, F, self.params) 22 23 def formJacobian(self, snes, X, J, P): 24 P.zeroEntries() # just in case 25 Bratu3D.FormJacobian(self.da, X, P, self.params) 26 if J != P: J.assemble() # matrix-free operator 27 return PETSc.Mat.Structure.SAME_NONZERO_PATTERN 28 29 30OptDB = PETSc.Options() 31 32N = OptDB.getInt('N', 16) 33lambda_ = OptDB.getReal('lambda', 6.0) 34do_plot = OptDB.getBool('plot', False) 35 36da = PETSc.DMDA().create([N, N, N], stencil_width=1) 37app = App(da, lambda_) 38 39snes = PETSc.SNES().create() 40F = da.createGlobalVec() 41snes.setFunction(app.formFunction, F) 42J = da.createMat() 43snes.setJacobian(app.formJacobian, J) 44 45snes.setFromOptions() 46 47X = da.createGlobalVec() 48app.formInitGuess(X) 49snes.solve(None, X) 50 51U = da.createNaturalVec() 52da.globalToNatural(X, U) 53 54 55def plot(da, U): 56 scatter, U0 = PETSc.Scatter.toZero(U) 57 scatter.scatter(U, U0, False, PETSc.Scatter.Mode.FORWARD) 58 rank = PETSc.COMM_WORLD.getRank() 59 if rank == 0: 60 solution = U0[...].reshape(da.sizes, order='f').copy() 61 try: 62 from matplotlib import pyplot 63 pyplot.contourf(solution[:, :, N//2]) 64 pyplot.axis('equal') 65 pyplot.show() 66 except: 67 pass 68 PETSc.COMM_WORLD.barrier() 69 scatter.destroy() 70 U0.destroy() 71 72 73if do_plot: plot(da, U) 74 75 76U.destroy() 77X.destroy() 78F.destroy() 79J.destroy() 80da.destroy() 81snes.destroy() 82