1# ------------------------------------------------------------------------
2#
3#  Solid Fuel Ignition (SFI) problem.  This problem is modeled by the
4#  partial differential equation
5#
6#          -Laplacian(u) - lambda * exp(u) = 0,  0 < x,y,z < 1,
7#
8#  with boundary conditions
9#
10#           u = 0  for  x = 0, x = 1, y = 0, y = 1, z = 0, z = 1
11#
12#  A finite difference approximation with the usual 7-point stencil
13#  is used to discretize the boundary value problem to obtain a
14#  nonlinear system of equations. The problem is solved in a 3D
15#  rectangular domain, using distributed arrays (DAs) to partition
16#  the parallel grid.
17#
18# ------------------------------------------------------------------------
19
20try: range = xrange
21except: pass
22
23import sys, petsc4py
24petsc4py.init(sys.argv)
25
26from numpy import exp, sqrt
27from petsc4py import PETSc
28
29class Bratu3D(object):
30
31    def __init__(self, da, lambda_):
32        assert da.getDim() == 3
33        self.da = da
34        self.lambda_ = lambda_
35        self.localX  = da.createLocalVec()
36
37    def formInitGuess(self, snes, X):
38        #
39        x = self.da.getVecArray(X)
40        #
41        mx, my, mz = self.da.getSizes()
42        hx, hy, hz = [1.0/(m-1) for m in [mx, my, mz]]
43        lambda_ = self.lambda_
44        scale = lambda_/(lambda_ + 1.0)
45        #
46        (xs, xe), (ys, ye), (zs, ze) = self.da.getRanges()
47        for k in range(zs, ze):
48            min_k = min(k,mz-k-1)*hz
49            for j in range(ys, ye):
50                min_j = min(j,my-j-1)*hy
51                for i in range(xs, xe):
52                    min_i = min(i,mx-i-1)*hx
53                    if (i==0    or j==0    or k==0 or
54                        i==mx-1 or j==my-1 or k==mz-1):
55                        # boundary points
56                        x[i, j, k] = 0.0
57                    else:
58                        # interior points
59                        min_kij = min(min_i,min_j,min_k)
60                        x[i, j, k] = scale*sqrt(min_kij)
61
62    def formFunction(self, snes, X, F):
63        #
64        self.da.globalToLocal(X, self.localX)
65        x = self.da.getVecArray(self.localX)
66        f = self.da.getVecArray(F)
67        #
68        mx, my, mz = self.da.getSizes()
69        hx, hy, hz = [1.0/m for m in [mx, my, mz]]
70        hxhyhz  = hx*hy*hz
71        hxhzdhy = hx*hz/hy;
72        hyhzdhx = hy*hz/hx;
73        hxhydhz = hx*hy/hz;
74        lambda_ = self.lambda_
75        #
76        (xs, xe), (ys, ye), (zs, ze) = self.da.getRanges()
77        for k in range(zs, ze):
78            for j in range(ys, ye):
79                for i in range(xs, xe):
80                    if (i==0    or j==0    or k==0 or
81                        i==mx-1 or j==my-1 or k==mz-1):
82                        f[i, j, k] = x[i, j, k] - 0
83                    else:
84                        u   = x[ i  ,  j  ,  k ] # center
85                        u_e = x[i+1 ,  j  ,  k ] # east
86                        u_w = x[i-1 ,  j  ,  k ] # west
87                        u_n = x[ i  , j+1 ,  k ] # north
88                        u_s = x[ i  , j-1 ,  k ] # south
89                        u_u = x[ i  ,  j  , k+1] # up
90                        u_d = x[ i  ,  j  , k-1] # down
91                        u_xx = (-u_e + 2*u - u_w)*hyhzdhx
92                        u_yy = (-u_n + 2*u - u_s)*hxhzdhy
93                        u_zz = (-u_u + 2*u - u_d)*hxhydhz
94                        f[i, j, k] = u_xx + u_yy + u_zz \
95                                     - lambda_*exp(u)*hxhyhz
96
97    def formJacobian(self, snes, X, J, P):
98        #
99        self.da.globalToLocal(X, self.localX)
100        x = self.da.getVecArray(self.localX)
101        #
102        mx, my, mz = self.da.getSizes()
103        hx, hy, hz = [1.0/m for m in [mx, my, mz]]
104        hxhyhz  = hx*hy*hz
105        hxhzdhy = hx*hz/hy;
106        hyhzdhx = hy*hz/hx;
107        hxhydhz = hx*hy/hz;
108        lambda_ = self.lambda_
109        #
110        P.zeroEntries()
111        row = PETSc.Mat.Stencil()
112        col = PETSc.Mat.Stencil()
113        #
114        (xs, xe), (ys, ye), (zs, ze) = self.da.getRanges()
115        for k in range(zs, ze):
116            for j in range(ys, ye):
117                for i in range(xs, xe):
118                    row.index = (i,j,k)
119                    row.field = 0
120                    if (i==0    or j==0    or k==0 or
121                        i==mx-1 or j==my-1 or k==mz-1):
122                        P.setValueStencil(row, row, 1.0)
123                    else:
124                        u = x[i,j,k]
125                        diag = (2*(hyhzdhx+hxhzdhy+hxhydhz)
126                                - lambda_*exp(u)*hxhyhz)
127                        for index, value in [
128                            ((i,j,k-1), -hxhydhz),
129                            ((i,j-1,k), -hxhzdhy),
130                            ((i-1,j,k), -hyhzdhx),
131                            ((i, j, k), diag),
132                            ((i+1,j,k), -hyhzdhx),
133                            ((i,j+1,k), -hxhzdhy),
134                            ((i,j,k+1), -hxhydhz),
135                            ]:
136                            col.index = index
137                            col.field = 0
138                            P.setValueStencil(row, col, value)
139        P.assemble()
140        if J != P: J.assemble() # matrix-free operator
141        return PETSc.Mat.Structure.SAME_NONZERO_PATTERN
142
143OptDB = PETSc.Options()
144
145n  = OptDB.getInt('n', 16)
146nx = OptDB.getInt('nx', n)
147ny = OptDB.getInt('ny', n)
148nz = OptDB.getInt('nz', n)
149lambda_ = OptDB.getReal('lambda', 6.0)
150
151da = PETSc.DMDA().create([nx, ny, nz], stencil_width=1)
152pde = Bratu3D(da, lambda_)
153
154snes = PETSc.SNES().create()
155F = da.createGlobalVec()
156snes.setFunction(pde.formFunction, F)
157
158fd = OptDB.getBool('fd', False)
159mf = OptDB.getBool('mf', False)
160if mf:
161    J = None
162    snes.setUseMF()
163else:
164    J = da.createMat()
165    snes.setJacobian(pde.formJacobian, J)
166    if fd:
167        snes.setUseFD()
168
169snes.getKSP().setType('cg')
170snes.setFromOptions()
171
172X = da.createGlobalVec()
173pde.formInitGuess(snes, X)
174snes.solve(None, X)
175
176U = da.createNaturalVec()
177da.globalToNatural(X, U)
178
179if OptDB.getBool('plot_mpl', False):
180
181    def plot_mpl(da, U):
182        comm = da.getComm()
183        rank = comm.getRank()
184        scatter, U0 = PETSc.Scatter.toZero(U)
185        scatter.scatter(U, U0, False, PETSc.Scatter.Mode.FORWARD)
186        if rank == 0:
187            try:
188                from matplotlib import pylab
189            except ImportError:
190                PETSc.Sys.Print("matplotlib not available")
191            else:
192                from numpy import mgrid
193                nx, ny, nz = da.sizes
194                solution = U0[...].reshape(da.sizes, order='f')
195                xx, yy =  mgrid[0:1:1j*nx,0:1:1j*ny]
196                pylab.contourf(xx, yy, solution[:, :, nz//2])
197                pylab.axis('equal')
198                pylab.xlabel('X')
199                pylab.ylabel('Y')
200                pylab.title('Z/2')
201                pylab.show()
202        comm.barrier()
203
204    plot_mpl(da, U)
205