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