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