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