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