xref: /petsc/src/dm/impls/stag/tutorials/ex3.c (revision 732aec7a18f2199fb53bb9a2f3aef439a834ce31)
1 static char help[] = "Solve a toy 3D problem on a staggered grid\n\n";
2 /*
3 
4   To demonstrate the basic functionality of DMStag, solves an isoviscous
5   incompressible Stokes problem on a rectangular 3D domain.
6 
7   u_{xx} + u_{yy} + u_{zz} - p_x = f^x
8   v_{xx} + v_{yy} + u_{zz} - p_y = f^y
9   w_{xx} + w_{yy} + w_{zz} - p_y = f^z
10   u_x    + v_y    + w_z          = g
11 
12   g = 0 for the physical case.
13 
14   Boundary conditions give prescribed flow perpendicular to the boundaries,
15   and zero derivative perpendicular to them (free slip). This involves
16   using a modified stencil at the boundaries. Another option would be to
17   use DM_BOUNDARY_GHOSTED in DMStagCreate3d() and a matrix-free operator (MATSHELL)
18   making use of the uniformly-available ghost layer.
19 
20   Use the -pinpressure option to fix a pressure node, instead of providing
21   a constant-pressure nullspace. This allows use of direct solvers, e.g. to
22   use UMFPACK,
23 
24      ./ex3 -pinpressure 1 -pc_type lu -pc_factor_mat_solver_type umfpack
25 
26 */
27 #include <petscdm.h>
28 #include <petscksp.h>
29 #include <petscdmstag.h>
30 
31 /* Shorter, more convenient names for DMStagStencilLocation entries */
32 #define BACK_DOWN_LEFT   DMSTAG_BACK_DOWN_LEFT
33 #define BACK_DOWN        DMSTAG_BACK_DOWN
34 #define BACK_DOWN_RIGHT  DMSTAG_BACK_DOWN_RIGHT
35 #define BACK_LEFT        DMSTAG_BACK_LEFT
36 #define BACK             DMSTAG_BACK
37 #define BACK_RIGHT       DMSTAG_BACK_RIGHT
38 #define BACK_UP_LEFT     DMSTAG_BACK_UP_LEFT
39 #define BACK_UP          DMSTAG_BACK_UP
40 #define BACK_UP_RIGHT    DMSTAG_BACK_UP_RIGHT
41 #define DOWN_LEFT        DMSTAG_DOWN_LEFT
42 #define DOWN             DMSTAG_DOWN
43 #define DOWN_RIGHT       DMSTAG_DOWN_RIGHT
44 #define LEFT             DMSTAG_LEFT
45 #define ELEMENT          DMSTAG_ELEMENT
46 #define RIGHT            DMSTAG_RIGHT
47 #define UP_LEFT          DMSTAG_UP_LEFT
48 #define UP               DMSTAG_UP
49 #define UP_RIGHT         DMSTAG_UP_RIGHT
50 #define FRONT_DOWN_LEFT  DMSTAG_FRONT_DOWN_LEFT
51 #define FRONT_DOWN       DMSTAG_FRONT_DOWN
52 #define FRONT_DOWN_RIGHT DMSTAG_FRONT_DOWN_RIGHT
53 #define FRONT_LEFT       DMSTAG_FRONT_LEFT
54 #define FRONT            DMSTAG_FRONT
55 #define FRONT_RIGHT      DMSTAG_FRONT_RIGHT
56 #define FRONT_UP_LEFT    DMSTAG_FRONT_UP_LEFT
57 #define FRONT_UP         DMSTAG_FRONT_UP
58 #define FRONT_UP_RIGHT   DMSTAG_FRONT_UP_RIGHT
59 
60 static PetscErrorCode CreateReferenceSolution(DM, Vec *);
61 static PetscErrorCode CreateSystem(DM, Mat *, Vec *, PetscBool);
62 static PetscErrorCode AttachNullspace(DM, Mat);
63 static PetscErrorCode CheckSolution(Vec, Vec);
64 
65 /* Manufactured solution. Chosen to be higher order than can be solved exactly,
66 and to have a zero derivative for flow parallel to the boundaries. That is,
67 d(ux)/dy = 0 at the top and bottom boundaries, and d(uy)/dx = 0 at the right
68 and left boundaries.
69 These expressions could be made more interesting with more z terms,
70 and convergence could be confirmed.  */
71 
uxRef(PetscScalar x,PetscScalar y,PetscScalar z)72 static PetscScalar uxRef(PetscScalar x, PetscScalar y, PetscScalar z)
73 {
74   return 0.0 * x + y * y - 2.0 * y * y * y + y * y * y * y + 0.0 * z;
75 }
uyRef(PetscScalar x,PetscScalar y,PetscScalar z)76 static PetscScalar uyRef(PetscScalar x, PetscScalar y, PetscScalar z)
77 {
78   return x * x - 2.0 * x * x * x + x * x * x * x + 0.0 * y + 0.0 * z;
79 }
uzRef(PetscScalar x,PetscScalar y,PetscScalar z)80 static PetscScalar uzRef(PetscScalar x, PetscScalar y, PetscScalar z)
81 {
82   return 0.0 * x + 0.0 * y + 0.0 * z + 1.0;
83 }
pRef(PetscScalar x,PetscScalar y,PetscScalar z)84 static PetscScalar pRef(PetscScalar x, PetscScalar y, PetscScalar z)
85 {
86   return -1.0 * (x - 0.5) + -3.0 / 2.0 * y * y + 0.5 - 2.0 * (z - 1.0);
87 } /* zero integral */
fx(PetscScalar x,PetscScalar y,PetscScalar z)88 static PetscScalar fx(PetscScalar x, PetscScalar y, PetscScalar z)
89 {
90   return 0.0 * x + 2.0 - 12.0 * y + 12.0 * y * y + 0.0 * z + 1.0;
91 }
fy(PetscScalar x,PetscScalar y,PetscScalar z)92 static PetscScalar fy(PetscScalar x, PetscScalar y, PetscScalar z)
93 {
94   return 2.0 - 12.0 * x + 12.0 * x * x + 3.0 * y + 0.0 * z;
95 }
fz(PetscScalar x,PetscScalar y,PetscScalar z)96 static PetscScalar fz(PetscScalar x, PetscScalar y, PetscScalar z)
97 {
98   return 0.0 * x + 0.0 * y + 0.0 * z + 2.0;
99 }
g(PetscScalar x,PetscScalar y,PetscScalar z)100 static PetscScalar g(PetscScalar x, PetscScalar y, PetscScalar z)
101 {
102   return 0.0 * x + 0.0 * y + 0.0 * z + 0.0;
103 }
104 
main(int argc,char ** argv)105 int main(int argc, char **argv)
106 {
107   DM        dmSol;
108   Vec       sol, solRef, rhs;
109   Mat       A;
110   KSP       ksp;
111   PC        pc;
112   PetscBool pinPressure;
113 
114   /* Initialize PETSc and process command line arguments */
115   PetscFunctionBeginUser;
116   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
117   pinPressure = PETSC_TRUE;
118   PetscCall(PetscOptionsGetBool(NULL, NULL, "-pinpressure", &pinPressure, NULL));
119 
120   /* Create 3D DMStag for the solution, and set up. */
121   {
122     const PetscInt dof0 = 0, dof1 = 0, dof2 = 1, dof3 = 1; /* 1 dof on each face and element center */
123     const PetscInt stencilWidth = 1;
124     PetscCall(DMStagCreate3d(PETSC_COMM_WORLD, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, 4, 5, 6, PETSC_DECIDE, PETSC_DECIDE, PETSC_DECIDE, dof0, dof1, dof2, dof3, DMSTAG_STENCIL_BOX, stencilWidth, NULL, NULL, NULL, &dmSol));
125     PetscCall(DMSetFromOptions(dmSol));
126     PetscCall(DMSetUp(dmSol));
127     PetscCall(DMStagSetUniformCoordinatesExplicit(dmSol, 0.0, 1.0, 0.0, 1.0, 0.0, 1.0));
128     /* Note: also see ex2.c, where another, more efficient option is demonstrated,
129        using DMStagSetUniformCoordinatesProduct() */
130   }
131 
132   /* Compute (manufactured) reference solution */
133   PetscCall(CreateReferenceSolution(dmSol, &solRef));
134 
135   /* Assemble system */
136   PetscCall(CreateSystem(dmSol, &A, &rhs, pinPressure));
137 
138   /* Attach a constant-pressure nullspace to the operator */
139   if (!pinPressure) PetscCall(AttachNullspace(dmSol, A));
140 
141   /* Solve */
142   PetscCall(DMCreateGlobalVector(dmSol, &sol));
143   PetscCall(KSPCreate(PETSC_COMM_WORLD, &ksp));
144   PetscCall(KSPSetType(ksp, KSPFGMRES));
145   PetscCall(KSPSetOperators(ksp, A, A));
146   PetscCall(KSPGetPC(ksp, &pc));
147   PetscCall(PCSetType(pc, PCFIELDSPLIT));
148   PetscCall(PCFieldSplitSetDetectSaddlePoint(pc, PETSC_TRUE));
149   PetscCall(KSPSetFromOptions(ksp));
150   PetscCall(KSPSolve(ksp, rhs, sol));
151 
152   /* Check Solution */
153   PetscCall(CheckSolution(sol, solRef));
154 
155   /* Clean up and finalize PETSc */
156   PetscCall(KSPDestroy(&ksp));
157   PetscCall(VecDestroy(&sol));
158   PetscCall(VecDestroy(&solRef));
159   PetscCall(VecDestroy(&rhs));
160   PetscCall(MatDestroy(&A));
161   PetscCall(DMDestroy(&dmSol));
162   PetscCall(PetscFinalize());
163   return 0;
164 }
165 
CreateSystem(DM dmSol,Mat * pA,Vec * pRhs,PetscBool pinPressure)166 static PetscErrorCode CreateSystem(DM dmSol, Mat *pA, Vec *pRhs, PetscBool pinPressure)
167 {
168   Vec             rhs, coordLocal;
169   Mat             A;
170   PetscInt        startx, starty, startz, N[3], nx, ny, nz, ex, ey, ez, d;
171   PetscInt        icp[3], icux[3], icuy[3], icuz[3], icux_right[3], icuy_up[3], icuz_front[3];
172   PetscReal       hx, hy, hz;
173   DM              dmCoord;
174   PetscScalar ****arrCoord;
175 
176   PetscFunctionBeginUser;
177   PetscCall(DMCreateMatrix(dmSol, pA));
178   A = *pA;
179   PetscCall(DMCreateGlobalVector(dmSol, pRhs));
180   rhs = *pRhs;
181 
182   PetscCall(DMStagGetCorners(dmSol, &startx, &starty, &startz, &nx, &ny, &nz, NULL, NULL, NULL));
183   PetscCall(DMStagGetGlobalSizes(dmSol, &N[0], &N[1], &N[2]));
184   PetscCheck(N[0] >= 2 && N[1] >= 2 && N[2] >= 2, PetscObjectComm((PetscObject)dmSol), PETSC_ERR_ARG_SIZ, "This example requires at least two elements in each dimensions");
185   hx = 1.0 / N[0];
186   hy = 1.0 / N[1];
187   hz = 1.0 / N[2];
188   PetscCall(DMGetCoordinateDM(dmSol, &dmCoord));
189   PetscCall(DMGetCoordinatesLocal(dmSol, &coordLocal));
190   PetscCall(DMStagVecGetArrayRead(dmCoord, coordLocal, &arrCoord));
191   for (d = 0; d < 3; ++d) {
192     PetscCall(DMStagGetLocationSlot(dmCoord, ELEMENT, d, &icp[d]));
193     PetscCall(DMStagGetLocationSlot(dmCoord, LEFT, d, &icux[d]));
194     PetscCall(DMStagGetLocationSlot(dmCoord, DOWN, d, &icuy[d]));
195     PetscCall(DMStagGetLocationSlot(dmCoord, BACK, d, &icuz[d]));
196     PetscCall(DMStagGetLocationSlot(dmCoord, RIGHT, d, &icux_right[d]));
197     PetscCall(DMStagGetLocationSlot(dmCoord, UP, d, &icuy_up[d]));
198     PetscCall(DMStagGetLocationSlot(dmCoord, FRONT, d, &icuz_front[d]));
199   }
200 
201   /* Loop over all local elements. Note that it may be more efficient in real
202      applications to loop over each boundary separately */
203   for (ez = startz; ez < startz + nz; ++ez) { /* With DMStag, always iterate x fastest, y second fastest, z slowest */
204     for (ey = starty; ey < starty + ny; ++ey) {
205       for (ex = startx; ex < startx + nx; ++ex) {
206         if (ex == N[0] - 1) {
207           /* Right Boundary velocity Dirichlet */
208           DMStagStencil     row;
209           PetscScalar       valRhs;
210           const PetscScalar valA = 1.0;
211           row.i                  = ex;
212           row.j                  = ey;
213           row.k                  = ez;
214           row.loc                = RIGHT;
215           row.c                  = 0;
216           PetscCall(DMStagMatSetValuesStencil(dmSol, A, 1, &row, 1, &row, &valA, INSERT_VALUES));
217           valRhs = uxRef(arrCoord[ez][ey][ex][icux_right[0]], arrCoord[ez][ey][ex][icux_right[1]], arrCoord[ez][ey][ex][icux_right[2]]);
218           PetscCall(DMStagVecSetValuesStencil(dmSol, rhs, 1, &row, &valRhs, INSERT_VALUES));
219         }
220         if (ey == N[1] - 1) {
221           /* Top boundary velocity Dirichlet */
222           DMStagStencil     row;
223           PetscScalar       valRhs;
224           const PetscScalar valA = 1.0;
225           row.i                  = ex;
226           row.j                  = ey;
227           row.k                  = ez;
228           row.loc                = UP;
229           row.c                  = 0;
230           PetscCall(DMStagMatSetValuesStencil(dmSol, A, 1, &row, 1, &row, &valA, INSERT_VALUES));
231           valRhs = uyRef(arrCoord[ez][ey][ex][icuy_up[0]], arrCoord[ez][ey][ex][icuy_up[1]], arrCoord[ez][ey][ex][icuy_up[2]]);
232           PetscCall(DMStagVecSetValuesStencil(dmSol, rhs, 1, &row, &valRhs, INSERT_VALUES));
233         }
234         if (ez == N[2] - 1) {
235           /* Front boundary velocity Dirichlet */
236           DMStagStencil     row;
237           PetscScalar       valRhs;
238           const PetscScalar valA = 1.0;
239           row.i                  = ex;
240           row.j                  = ey;
241           row.k                  = ez;
242           row.loc                = FRONT;
243           row.c                  = 0;
244           PetscCall(DMStagMatSetValuesStencil(dmSol, A, 1, &row, 1, &row, &valA, INSERT_VALUES));
245           valRhs = uzRef(arrCoord[ez][ey][ex][icuz_front[0]], arrCoord[ez][ey][ex][icuz_front[1]], arrCoord[ez][ey][ex][icuz_front[2]]);
246           PetscCall(DMStagVecSetValuesStencil(dmSol, rhs, 1, &row, &valRhs, INSERT_VALUES));
247         }
248 
249         /* Equation on left face of this element */
250         if (ex == 0) {
251           /* Left velocity Dirichlet */
252           DMStagStencil     row;
253           PetscScalar       valRhs;
254           const PetscScalar valA = 1.0;
255           row.i                  = ex;
256           row.j                  = ey;
257           row.k                  = ez;
258           row.loc                = LEFT;
259           row.c                  = 0;
260           PetscCall(DMStagMatSetValuesStencil(dmSol, A, 1, &row, 1, &row, &valA, INSERT_VALUES));
261           valRhs = uxRef(arrCoord[ez][ey][ex][icux[0]], arrCoord[ez][ey][ex][icux[1]], arrCoord[ez][ey][ex][icux[2]]);
262           PetscCall(DMStagVecSetValuesStencil(dmSol, rhs, 1, &row, &valRhs, INSERT_VALUES));
263         } else {
264           /* X-momentum interior equation : (u_xx + u_yy + u_zz) - p_x = f^x */
265           DMStagStencil row, col[9];
266           PetscScalar   valA[9], valRhs;
267           PetscInt      nEntries;
268 
269           row.i   = ex;
270           row.j   = ey;
271           row.k   = ez;
272           row.loc = LEFT;
273           row.c   = 0;
274           if (ey == 0) {
275             if (ez == 0) {
276               nEntries   = 7;
277               col[0].i   = ex;
278               col[0].j   = ey;
279               col[0].k   = ez;
280               col[0].loc = LEFT;
281               col[0].c   = 0;
282               valA[0]    = -2.0 / (hx * hx) + -1.0 / (hy * hy) - 1.0 / (hz * hz);
283               /* Missing down term */
284               col[1].i   = ex;
285               col[1].j   = ey + 1;
286               col[1].k   = ez;
287               col[1].loc = LEFT;
288               col[1].c   = 0;
289               valA[1]    = 1.0 / (hy * hy);
290               col[2].i   = ex - 1;
291               col[2].j   = ey;
292               col[2].k   = ez;
293               col[2].loc = LEFT;
294               col[2].c   = 0;
295               valA[2]    = 1.0 / (hx * hx);
296               col[3].i   = ex + 1;
297               col[3].j   = ey;
298               col[3].k   = ez;
299               col[3].loc = LEFT;
300               col[3].c   = 0;
301               valA[3]    = 1.0 / (hx * hx);
302               /* Missing back term */
303               col[4].i   = ex;
304               col[4].j   = ey;
305               col[4].k   = ez + 1;
306               col[4].loc = LEFT;
307               col[4].c   = 0;
308               valA[4]    = 1.0 / (hz * hz);
309               col[5].i   = ex - 1;
310               col[5].j   = ey;
311               col[5].k   = ez;
312               col[5].loc = ELEMENT;
313               col[5].c   = 0;
314               valA[5]    = 1.0 / hx;
315               col[6].i   = ex;
316               col[6].j   = ey;
317               col[6].k   = ez;
318               col[6].loc = ELEMENT;
319               col[6].c   = 0;
320               valA[6]    = -1.0 / hx;
321             } else if (ez == N[2] - 1) {
322               nEntries   = 7;
323               col[0].i   = ex;
324               col[0].j   = ey;
325               col[0].k   = ez;
326               col[0].loc = LEFT;
327               col[0].c   = 0;
328               valA[0]    = -2.0 / (hx * hx) + -1.0 / (hy * hy) - 1.0 / (hz * hz);
329               /* Missing down term */
330               col[1].i   = ex;
331               col[1].j   = ey + 1;
332               col[1].k   = ez;
333               col[1].loc = LEFT;
334               col[1].c   = 0;
335               valA[1]    = 1.0 / (hy * hy);
336               col[2].i   = ex - 1;
337               col[2].j   = ey;
338               col[2].k   = ez;
339               col[2].loc = LEFT;
340               col[2].c   = 0;
341               valA[2]    = 1.0 / (hx * hx);
342               col[3].i   = ex + 1;
343               col[3].j   = ey;
344               col[3].k   = ez;
345               col[3].loc = LEFT;
346               col[3].c   = 0;
347               valA[3]    = 1.0 / (hx * hx);
348               col[4].i   = ex;
349               col[4].j   = ey;
350               col[4].k   = ez - 1;
351               col[4].loc = LEFT;
352               col[4].c   = 0;
353               valA[4]    = 1.0 / (hz * hz);
354               /* Missing front term */
355               col[5].i   = ex - 1;
356               col[5].j   = ey;
357               col[5].k   = ez;
358               col[5].loc = ELEMENT;
359               col[5].c   = 0;
360               valA[5]    = 1.0 / hx;
361               col[6].i   = ex;
362               col[6].j   = ey;
363               col[6].k   = ez;
364               col[6].loc = ELEMENT;
365               col[6].c   = 0;
366               valA[6]    = -1.0 / hx;
367             } else {
368               nEntries   = 8;
369               col[0].i   = ex;
370               col[0].j   = ey;
371               col[0].k   = ez;
372               col[0].loc = LEFT;
373               col[0].c   = 0;
374               valA[0]    = -2.0 / (hx * hx) + -1.0 / (hy * hy) - 2.0 / (hz * hz);
375               /* Missing down term */
376               col[1].i   = ex;
377               col[1].j   = ey + 1;
378               col[1].k   = ez;
379               col[1].loc = LEFT;
380               col[1].c   = 0;
381               valA[1]    = 1.0 / (hy * hy);
382               col[2].i   = ex - 1;
383               col[2].j   = ey;
384               col[2].k   = ez;
385               col[2].loc = LEFT;
386               col[2].c   = 0;
387               valA[2]    = 1.0 / (hx * hx);
388               col[3].i   = ex + 1;
389               col[3].j   = ey;
390               col[3].k   = ez;
391               col[3].loc = LEFT;
392               col[3].c   = 0;
393               valA[3]    = 1.0 / (hx * hx);
394               col[4].i   = ex;
395               col[4].j   = ey;
396               col[4].k   = ez - 1;
397               col[4].loc = LEFT;
398               col[4].c   = 0;
399               valA[4]    = 1.0 / (hz * hz);
400               col[5].i   = ex;
401               col[5].j   = ey;
402               col[5].k   = ez + 1;
403               col[5].loc = LEFT;
404               col[5].c   = 0;
405               valA[5]    = 1.0 / (hz * hz);
406               col[6].i   = ex - 1;
407               col[6].j   = ey;
408               col[6].k   = ez;
409               col[6].loc = ELEMENT;
410               col[6].c   = 0;
411               valA[6]    = 1.0 / hx;
412               col[7].i   = ex;
413               col[7].j   = ey;
414               col[7].k   = ez;
415               col[7].loc = ELEMENT;
416               col[7].c   = 0;
417               valA[7]    = -1.0 / hx;
418             }
419           } else if (ey == N[1] - 1) {
420             if (ez == 0) {
421               nEntries   = 7;
422               col[0].i   = ex;
423               col[0].j   = ey;
424               col[0].k   = ez;
425               col[0].loc = LEFT;
426               col[0].c   = 0;
427               valA[0]    = -2.0 / (hx * hx) + -2.0 / (hy * hy) - 1.0 / (hz * hz);
428               col[1].i   = ex;
429               col[1].j   = ey - 1;
430               col[1].k   = ez;
431               col[1].loc = LEFT;
432               col[1].c   = 0;
433               valA[1]    = 1.0 / (hy * hy);
434               /* Missing up term */
435               col[2].i   = ex - 1;
436               col[2].j   = ey;
437               col[2].k   = ez;
438               col[2].loc = LEFT;
439               col[2].c   = 0;
440               valA[2]    = 1.0 / (hx * hx);
441               col[3].i   = ex + 1;
442               col[3].j   = ey;
443               col[3].k   = ez;
444               col[3].loc = LEFT;
445               col[3].c   = 0;
446               valA[3]    = 1.0 / (hx * hx);
447               /* Missing back entry */
448               col[4].i   = ex;
449               col[4].j   = ey;
450               col[4].k   = ez + 1;
451               col[4].loc = LEFT;
452               col[4].c   = 0;
453               valA[4]    = 1.0 / (hz * hz);
454               col[5].i   = ex - 1;
455               col[5].j   = ey;
456               col[5].k   = ez;
457               col[5].loc = ELEMENT;
458               col[5].c   = 0;
459               valA[5]    = 1.0 / hx;
460               col[6].i   = ex;
461               col[6].j   = ey;
462               col[6].k   = ez;
463               col[6].loc = ELEMENT;
464               col[6].c   = 0;
465               valA[6]    = -1.0 / hx;
466             } else if (ez == N[2] - 1) {
467               nEntries   = 7;
468               col[0].i   = ex;
469               col[0].j   = ey;
470               col[0].k   = ez;
471               col[0].loc = LEFT;
472               col[0].c   = 0;
473               valA[0]    = -2.0 / (hx * hx) + -2.0 / (hy * hy) - 1.0 / (hz * hz);
474               col[1].i   = ex;
475               col[1].j   = ey - 1;
476               col[1].k   = ez;
477               col[1].loc = LEFT;
478               col[1].c   = 0;
479               valA[1]    = 1.0 / (hy * hy);
480               /* Missing up term */
481               col[2].i   = ex - 1;
482               col[2].j   = ey;
483               col[2].k   = ez;
484               col[2].loc = LEFT;
485               col[2].c   = 0;
486               valA[2]    = 1.0 / (hx * hx);
487               col[3].i   = ex + 1;
488               col[3].j   = ey;
489               col[3].k   = ez;
490               col[3].loc = LEFT;
491               col[3].c   = 0;
492               valA[3]    = 1.0 / (hx * hx);
493               col[4].i   = ex;
494               col[4].j   = ey;
495               col[4].k   = ez - 1;
496               col[4].loc = LEFT;
497               col[4].c   = 0;
498               valA[4]    = 1.0 / (hz * hz);
499               /* Missing front term */
500               col[5].i   = ex - 1;
501               col[5].j   = ey;
502               col[5].k   = ez;
503               col[5].loc = ELEMENT;
504               col[5].c   = 0;
505               valA[5]    = 1.0 / hx;
506               col[6].i   = ex;
507               col[6].j   = ey;
508               col[6].k   = ez;
509               col[6].loc = ELEMENT;
510               col[6].c   = 0;
511               valA[6]    = -1.0 / hx;
512             } else {
513               nEntries   = 8;
514               col[0].i   = ex;
515               col[0].j   = ey;
516               col[0].k   = ez;
517               col[0].loc = LEFT;
518               col[0].c   = 0;
519               valA[0]    = -2.0 / (hx * hx) + -2.0 / (hy * hy) - 2.0 / (hz * hz);
520               col[1].i   = ex;
521               col[1].j   = ey - 1;
522               col[1].k   = ez;
523               col[1].loc = LEFT;
524               col[1].c   = 0;
525               valA[1]    = 1.0 / (hy * hy);
526               /* Missing up term */
527               col[2].i   = ex - 1;
528               col[2].j   = ey;
529               col[2].k   = ez;
530               col[2].loc = LEFT;
531               col[2].c   = 0;
532               valA[2]    = 1.0 / (hx * hx);
533               col[3].i   = ex + 1;
534               col[3].j   = ey;
535               col[3].k   = ez;
536               col[3].loc = LEFT;
537               col[3].c   = 0;
538               valA[3]    = 1.0 / (hx * hx);
539               col[4].i   = ex;
540               col[4].j   = ey;
541               col[4].k   = ez - 1;
542               col[4].loc = LEFT;
543               col[4].c   = 0;
544               valA[4]    = 1.0 / (hz * hz);
545               col[5].i   = ex;
546               col[5].j   = ey;
547               col[5].k   = ez + 1;
548               col[5].loc = LEFT;
549               col[5].c   = 0;
550               valA[5]    = 1.0 / (hz * hz);
551               col[6].i   = ex - 1;
552               col[6].j   = ey;
553               col[6].k   = ez;
554               col[6].loc = ELEMENT;
555               col[6].c   = 0;
556               valA[6]    = 1.0 / hx;
557               col[7].i   = ex;
558               col[7].j   = ey;
559               col[7].k   = ez;
560               col[7].loc = ELEMENT;
561               col[7].c   = 0;
562               valA[7]    = -1.0 / hx;
563             }
564           } else if (ez == 0) {
565             nEntries   = 8;
566             col[0].i   = ex;
567             col[0].j   = ey;
568             col[0].k   = ez;
569             col[0].loc = LEFT;
570             col[0].c   = 0;
571             valA[0]    = -2.0 / (hx * hx) + -2.0 / (hy * hy) - 1.0 / (hz * hz);
572             col[1].i   = ex;
573             col[1].j   = ey - 1;
574             col[1].k   = ez;
575             col[1].loc = LEFT;
576             col[1].c   = 0;
577             valA[1]    = 1.0 / (hy * hy);
578             col[2].i   = ex;
579             col[2].j   = ey + 1;
580             col[2].k   = ez;
581             col[2].loc = LEFT;
582             col[2].c   = 0;
583             valA[2]    = 1.0 / (hy * hy);
584             col[3].i   = ex - 1;
585             col[3].j   = ey;
586             col[3].k   = ez;
587             col[3].loc = LEFT;
588             col[3].c   = 0;
589             valA[3]    = 1.0 / (hx * hx);
590             col[4].i   = ex + 1;
591             col[4].j   = ey;
592             col[4].k   = ez;
593             col[4].loc = LEFT;
594             col[4].c   = 0;
595             valA[4]    = 1.0 / (hx * hx);
596             /* Missing back term */
597             col[5].i   = ex;
598             col[5].j   = ey;
599             col[5].k   = ez + 1;
600             col[5].loc = LEFT;
601             col[5].c   = 0;
602             valA[5]    = 1.0 / (hz * hz);
603             col[6].i   = ex - 1;
604             col[6].j   = ey;
605             col[6].k   = ez;
606             col[6].loc = ELEMENT;
607             col[6].c   = 0;
608             valA[6]    = 1.0 / hx;
609             col[7].i   = ex;
610             col[7].j   = ey;
611             col[7].k   = ez;
612             col[7].loc = ELEMENT;
613             col[7].c   = 0;
614             valA[7]    = -1.0 / hx;
615           } else if (ez == N[2] - 1) {
616             nEntries   = 8;
617             col[0].i   = ex;
618             col[0].j   = ey;
619             col[0].k   = ez;
620             col[0].loc = LEFT;
621             col[0].c   = 0;
622             valA[0]    = -2.0 / (hx * hx) + -2.0 / (hy * hy) - 1.0 / (hz * hz);
623             col[1].i   = ex;
624             col[1].j   = ey - 1;
625             col[1].k   = ez;
626             col[1].loc = LEFT;
627             col[1].c   = 0;
628             valA[1]    = 1.0 / (hy * hy);
629             col[2].i   = ex;
630             col[2].j   = ey + 1;
631             col[2].k   = ez;
632             col[2].loc = LEFT;
633             col[2].c   = 0;
634             valA[2]    = 1.0 / (hy * hy);
635             col[3].i   = ex - 1;
636             col[3].j   = ey;
637             col[3].k   = ez;
638             col[3].loc = LEFT;
639             col[3].c   = 0;
640             valA[3]    = 1.0 / (hx * hx);
641             col[4].i   = ex + 1;
642             col[4].j   = ey;
643             col[4].k   = ez;
644             col[4].loc = LEFT;
645             col[4].c   = 0;
646             valA[4]    = 1.0 / (hx * hx);
647             col[5].i   = ex;
648             col[5].j   = ey;
649             col[5].k   = ez - 1;
650             col[5].loc = LEFT;
651             col[5].c   = 0;
652             valA[5]    = 1.0 / (hz * hz);
653             /* Missing front term */
654             col[6].i   = ex - 1;
655             col[6].j   = ey;
656             col[6].k   = ez;
657             col[6].loc = ELEMENT;
658             col[6].c   = 0;
659             valA[6]    = 1.0 / hx;
660             col[7].i   = ex;
661             col[7].j   = ey;
662             col[7].k   = ez;
663             col[7].loc = ELEMENT;
664             col[7].c   = 0;
665             valA[7]    = -1.0 / hx;
666           } else {
667             nEntries   = 9;
668             col[0].i   = ex;
669             col[0].j   = ey;
670             col[0].k   = ez;
671             col[0].loc = LEFT;
672             col[0].c   = 0;
673             valA[0]    = -2.0 / (hx * hx) + -2.0 / (hy * hy) - 2.0 / (hz * hz);
674             col[1].i   = ex;
675             col[1].j   = ey - 1;
676             col[1].k   = ez;
677             col[1].loc = LEFT;
678             col[1].c   = 0;
679             valA[1]    = 1.0 / (hy * hy);
680             col[2].i   = ex;
681             col[2].j   = ey + 1;
682             col[2].k   = ez;
683             col[2].loc = LEFT;
684             col[2].c   = 0;
685             valA[2]    = 1.0 / (hy * hy);
686             col[3].i   = ex - 1;
687             col[3].j   = ey;
688             col[3].k   = ez;
689             col[3].loc = LEFT;
690             col[3].c   = 0;
691             valA[3]    = 1.0 / (hx * hx);
692             col[4].i   = ex + 1;
693             col[4].j   = ey;
694             col[4].k   = ez;
695             col[4].loc = LEFT;
696             col[4].c   = 0;
697             valA[4]    = 1.0 / (hx * hx);
698             col[5].i   = ex;
699             col[5].j   = ey;
700             col[5].k   = ez - 1;
701             col[5].loc = LEFT;
702             col[5].c   = 0;
703             valA[5]    = 1.0 / (hz * hz);
704             col[6].i   = ex;
705             col[6].j   = ey;
706             col[6].k   = ez + 1;
707             col[6].loc = LEFT;
708             col[6].c   = 0;
709             valA[6]    = 1.0 / (hz * hz);
710             col[7].i   = ex - 1;
711             col[7].j   = ey;
712             col[7].k   = ez;
713             col[7].loc = ELEMENT;
714             col[7].c   = 0;
715             valA[7]    = 1.0 / hx;
716             col[8].i   = ex;
717             col[8].j   = ey;
718             col[8].k   = ez;
719             col[8].loc = ELEMENT;
720             col[8].c   = 0;
721             valA[8]    = -1.0 / hx;
722           }
723           PetscCall(DMStagMatSetValuesStencil(dmSol, A, 1, &row, nEntries, col, valA, INSERT_VALUES));
724           valRhs = fx(arrCoord[ez][ey][ex][icux[0]], arrCoord[ez][ey][ex][icux[1]], arrCoord[ez][ey][ex][icux[2]]);
725           PetscCall(DMStagVecSetValuesStencil(dmSol, rhs, 1, &row, &valRhs, INSERT_VALUES));
726         }
727 
728         /* Equation on bottom face of this element */
729         if (ey == 0) {
730           /* Bottom boundary velocity Dirichlet */
731           DMStagStencil     row;
732           PetscScalar       valRhs;
733           const PetscScalar valA = 1.0;
734           row.i                  = ex;
735           row.j                  = ey;
736           row.k                  = ez;
737           row.loc                = DOWN;
738           row.c                  = 0;
739           PetscCall(DMStagMatSetValuesStencil(dmSol, A, 1, &row, 1, &row, &valA, INSERT_VALUES));
740           valRhs = uyRef(arrCoord[ez][ey][ex][icuy[0]], arrCoord[ez][ey][ex][icuy[1]], arrCoord[ez][ey][ex][icuy[2]]);
741           PetscCall(DMStagVecSetValuesStencil(dmSol, rhs, 1, &row, &valRhs, INSERT_VALUES));
742         } else {
743           /* Y-momentum equation, (v_xx + v_yy + v_zz) - p_y = f^y */
744           DMStagStencil row, col[9];
745           PetscScalar   valA[9], valRhs;
746           PetscInt      nEntries;
747 
748           row.i   = ex;
749           row.j   = ey;
750           row.k   = ez;
751           row.loc = DOWN;
752           row.c   = 0;
753           if (ex == 0) {
754             if (ez == 0) {
755               nEntries   = 7;
756               col[0].i   = ex;
757               col[0].j   = ey;
758               col[0].k   = ez;
759               col[0].loc = DOWN;
760               col[0].c   = 0;
761               valA[0]    = -1.0 / (hx * hx) + -2.0 / (hy * hy) - 1.0 / (hz * hz);
762               col[1].i   = ex;
763               col[1].j   = ey - 1;
764               col[1].k   = ez;
765               col[1].loc = DOWN;
766               col[1].c   = 0;
767               valA[1]    = 1.0 / (hy * hy);
768               col[2].i   = ex;
769               col[2].j   = ey + 1;
770               col[2].k   = ez;
771               col[2].loc = DOWN;
772               col[2].c   = 0;
773               valA[2]    = 1.0 / (hy * hy);
774               /* Left term missing */
775               col[3].i   = ex + 1;
776               col[3].j   = ey;
777               col[3].k   = ez;
778               col[3].loc = DOWN;
779               col[3].c   = 0;
780               valA[3]    = 1.0 / (hx * hx);
781               /* Back term missing */
782               col[4].i   = ex;
783               col[4].j   = ey;
784               col[4].k   = ez + 1;
785               col[4].loc = DOWN;
786               col[4].c   = 0;
787               valA[4]    = 1.0 / (hz * hz);
788               col[5].i   = ex;
789               col[5].j   = ey - 1;
790               col[5].k   = ez;
791               col[5].loc = ELEMENT;
792               col[5].c   = 0;
793               valA[5]    = 1.0 / hy;
794               col[6].i   = ex;
795               col[6].j   = ey;
796               col[6].k   = ez;
797               col[6].loc = ELEMENT;
798               col[6].c   = 0;
799               valA[6]    = -1.0 / hy;
800             } else if (ez == N[2] - 1) {
801               nEntries   = 7;
802               col[0].i   = ex;
803               col[0].j   = ey;
804               col[0].k   = ez;
805               col[0].loc = DOWN;
806               col[0].c   = 0;
807               valA[0]    = -1.0 / (hx * hx) + -2.0 / (hy * hy) - 1.0 / (hz * hz);
808               col[1].i   = ex;
809               col[1].j   = ey - 1;
810               col[1].k   = ez;
811               col[1].loc = DOWN;
812               col[1].c   = 0;
813               valA[1]    = 1.0 / (hy * hy);
814               col[2].i   = ex;
815               col[2].j   = ey + 1;
816               col[2].k   = ez;
817               col[2].loc = DOWN;
818               col[2].c   = 0;
819               valA[2]    = 1.0 / (hy * hy);
820               /* Left term missing */
821               col[3].i   = ex + 1;
822               col[3].j   = ey;
823               col[3].k   = ez;
824               col[3].loc = DOWN;
825               col[3].c   = 0;
826               valA[3]    = 1.0 / (hx * hx);
827               col[4].i   = ex;
828               col[4].j   = ey;
829               col[4].k   = ez - 1;
830               col[4].loc = DOWN;
831               col[4].c   = 0;
832               valA[4]    = 1.0 / (hz * hz);
833               /* Front term missing */
834               col[5].i   = ex;
835               col[5].j   = ey - 1;
836               col[5].k   = ez;
837               col[5].loc = ELEMENT;
838               col[5].c   = 0;
839               valA[5]    = 1.0 / hy;
840               col[6].i   = ex;
841               col[6].j   = ey;
842               col[6].k   = ez;
843               col[6].loc = ELEMENT;
844               col[6].c   = 0;
845               valA[6]    = -1.0 / hy;
846             } else {
847               nEntries   = 8;
848               col[0].i   = ex;
849               col[0].j   = ey;
850               col[0].k   = ez;
851               col[0].loc = DOWN;
852               col[0].c   = 0;
853               valA[0]    = -1.0 / (hx * hx) + -2.0 / (hy * hy) - 2.0 / (hz * hz);
854               col[1].i   = ex;
855               col[1].j   = ey - 1;
856               col[1].k   = ez;
857               col[1].loc = DOWN;
858               col[1].c   = 0;
859               valA[1]    = 1.0 / (hy * hy);
860               col[2].i   = ex;
861               col[2].j   = ey + 1;
862               col[2].k   = ez;
863               col[2].loc = DOWN;
864               col[2].c   = 0;
865               valA[2]    = 1.0 / (hy * hy);
866               /* Left term missing */
867               col[3].i   = ex + 1;
868               col[3].j   = ey;
869               col[3].k   = ez;
870               col[3].loc = DOWN;
871               col[3].c   = 0;
872               valA[3]    = 1.0 / (hx * hx);
873               col[4].i   = ex;
874               col[4].j   = ey;
875               col[4].k   = ez - 1;
876               col[4].loc = DOWN;
877               col[4].c   = 0;
878               valA[4]    = 1.0 / (hz * hz);
879               col[5].i   = ex;
880               col[5].j   = ey;
881               col[5].k   = ez + 1;
882               col[5].loc = DOWN;
883               col[5].c   = 0;
884               valA[5]    = 1.0 / (hz * hz);
885               col[6].i   = ex;
886               col[6].j   = ey - 1;
887               col[6].k   = ez;
888               col[6].loc = ELEMENT;
889               col[6].c   = 0;
890               valA[6]    = 1.0 / hy;
891               col[7].i   = ex;
892               col[7].j   = ey;
893               col[7].k   = ez;
894               col[7].loc = ELEMENT;
895               col[7].c   = 0;
896               valA[7]    = -1.0 / hy;
897             }
898           } else if (ex == N[0] - 1) {
899             if (ez == 0) {
900               nEntries   = 7;
901               col[0].i   = ex;
902               col[0].j   = ey;
903               col[0].k   = ez;
904               col[0].loc = DOWN;
905               col[0].c   = 0;
906               valA[0]    = -1.0 / (hx * hx) + -2.0 / (hy * hy) - 1.0 / (hz * hz);
907               col[1].i   = ex;
908               col[1].j   = ey - 1;
909               col[1].k   = ez;
910               col[1].loc = DOWN;
911               col[1].c   = 0;
912               valA[1]    = 1.0 / (hy * hy);
913               col[2].i   = ex;
914               col[2].j   = ey + 1;
915               col[2].k   = ez;
916               col[2].loc = DOWN;
917               col[2].c   = 0;
918               valA[2]    = 1.0 / (hy * hy);
919               col[3].i   = ex - 1;
920               col[3].j   = ey;
921               col[3].k   = ez;
922               col[3].loc = DOWN;
923               col[3].c   = 0;
924               valA[3]    = 1.0 / (hx * hx);
925               /* Right term missing */
926               /* Back term missing */
927               col[4].i   = ex;
928               col[4].j   = ey;
929               col[4].k   = ez + 1;
930               col[4].loc = DOWN;
931               col[4].c   = 0;
932               valA[4]    = 1.0 / (hz * hz);
933               col[5].i   = ex;
934               col[5].j   = ey - 1;
935               col[5].k   = ez;
936               col[5].loc = ELEMENT;
937               col[5].c   = 0;
938               valA[5]    = 1.0 / hy;
939               col[6].i   = ex;
940               col[6].j   = ey;
941               col[6].k   = ez;
942               col[6].loc = ELEMENT;
943               col[6].c   = 0;
944               valA[6]    = -1.0 / hy;
945             } else if (ez == N[2] - 1) {
946               nEntries   = 7;
947               col[0].i   = ex;
948               col[0].j   = ey;
949               col[0].k   = ez;
950               col[0].loc = DOWN;
951               col[0].c   = 0;
952               valA[0]    = -1.0 / (hx * hx) + -2.0 / (hy * hy) - 1.0 / (hz * hz);
953               col[1].i   = ex;
954               col[1].j   = ey - 1;
955               col[1].k   = ez;
956               col[1].loc = DOWN;
957               col[1].c   = 0;
958               valA[1]    = 1.0 / (hy * hy);
959               col[2].i   = ex;
960               col[2].j   = ey + 1;
961               col[2].k   = ez;
962               col[2].loc = DOWN;
963               col[2].c   = 0;
964               valA[2]    = 1.0 / (hy * hy);
965               col[3].i   = ex - 1;
966               col[3].j   = ey;
967               col[3].k   = ez;
968               col[3].loc = DOWN;
969               col[3].c   = 0;
970               valA[3]    = 1.0 / (hx * hx);
971               /* Right term missing */
972               col[4].i   = ex;
973               col[4].j   = ey;
974               col[4].k   = ez - 1;
975               col[4].loc = DOWN;
976               col[4].c   = 0;
977               valA[4]    = 1.0 / (hz * hz);
978               /* Front term missing */
979               col[5].i   = ex;
980               col[5].j   = ey - 1;
981               col[5].k   = ez;
982               col[5].loc = ELEMENT;
983               col[5].c   = 0;
984               valA[5]    = 1.0 / hy;
985               col[6].i   = ex;
986               col[6].j   = ey;
987               col[6].k   = ez;
988               col[6].loc = ELEMENT;
989               col[6].c   = 0;
990               valA[6]    = -1.0 / hy;
991             } else {
992               nEntries   = 8;
993               col[0].i   = ex;
994               col[0].j   = ey;
995               col[0].k   = ez;
996               col[0].loc = DOWN;
997               col[0].c   = 0;
998               valA[0]    = -1.0 / (hx * hx) + -2.0 / (hy * hy) - 2.0 / (hz * hz);
999               col[1].i   = ex;
1000               col[1].j   = ey - 1;
1001               col[1].k   = ez;
1002               col[1].loc = DOWN;
1003               col[1].c   = 0;
1004               valA[1]    = 1.0 / (hy * hy);
1005               col[2].i   = ex;
1006               col[2].j   = ey + 1;
1007               col[2].k   = ez;
1008               col[2].loc = DOWN;
1009               col[2].c   = 0;
1010               valA[2]    = 1.0 / (hy * hy);
1011               col[3].i   = ex - 1;
1012               col[3].j   = ey;
1013               col[3].k   = ez;
1014               col[3].loc = DOWN;
1015               col[3].c   = 0;
1016               valA[3]    = 1.0 / (hx * hx);
1017               /* Right term missing */
1018               col[4].i   = ex;
1019               col[4].j   = ey;
1020               col[4].k   = ez - 1;
1021               col[4].loc = DOWN;
1022               col[4].c   = 0;
1023               valA[4]    = 1.0 / (hz * hz);
1024               col[5].i   = ex;
1025               col[5].j   = ey;
1026               col[5].k   = ez + 1;
1027               col[5].loc = DOWN;
1028               col[5].c   = 0;
1029               valA[5]    = 1.0 / (hz * hz);
1030               col[6].i   = ex;
1031               col[6].j   = ey - 1;
1032               col[6].k   = ez;
1033               col[6].loc = ELEMENT;
1034               col[6].c   = 0;
1035               valA[6]    = 1.0 / hy;
1036               col[7].i   = ex;
1037               col[7].j   = ey;
1038               col[7].k   = ez;
1039               col[7].loc = ELEMENT;
1040               col[7].c   = 0;
1041               valA[7]    = -1.0 / hy;
1042             }
1043           } else if (ez == 0) {
1044             nEntries   = 8;
1045             col[0].i   = ex;
1046             col[0].j   = ey;
1047             col[0].k   = ez;
1048             col[0].loc = DOWN;
1049             col[0].c   = 0;
1050             valA[0]    = -2.0 / (hx * hx) + -2.0 / (hy * hy) - 1.0 / (hz * hz);
1051             col[1].i   = ex;
1052             col[1].j   = ey - 1;
1053             col[1].k   = ez;
1054             col[1].loc = DOWN;
1055             col[1].c   = 0;
1056             valA[1]    = 1.0 / (hy * hy);
1057             col[2].i   = ex;
1058             col[2].j   = ey + 1;
1059             col[2].k   = ez;
1060             col[2].loc = DOWN;
1061             col[2].c   = 0;
1062             valA[2]    = 1.0 / (hy * hy);
1063             col[3].i   = ex - 1;
1064             col[3].j   = ey;
1065             col[3].k   = ez;
1066             col[3].loc = DOWN;
1067             col[3].c   = 0;
1068             valA[3]    = 1.0 / (hx * hx);
1069             col[4].i   = ex + 1;
1070             col[4].j   = ey;
1071             col[4].k   = ez;
1072             col[4].loc = DOWN;
1073             col[4].c   = 0;
1074             valA[4]    = 1.0 / (hx * hx);
1075             /* Back term missing */
1076             col[5].i   = ex;
1077             col[5].j   = ey;
1078             col[5].k   = ez + 1;
1079             col[5].loc = DOWN;
1080             col[5].c   = 0;
1081             valA[5]    = 1.0 / (hz * hz);
1082             col[6].i   = ex;
1083             col[6].j   = ey - 1;
1084             col[6].k   = ez;
1085             col[6].loc = ELEMENT;
1086             col[6].c   = 0;
1087             valA[6]    = 1.0 / hy;
1088             col[7].i   = ex;
1089             col[7].j   = ey;
1090             col[7].k   = ez;
1091             col[7].loc = ELEMENT;
1092             col[7].c   = 0;
1093             valA[7]    = -1.0 / hy;
1094           } else if (ez == N[2] - 1) {
1095             nEntries   = 8;
1096             col[0].i   = ex;
1097             col[0].j   = ey;
1098             col[0].k   = ez;
1099             col[0].loc = DOWN;
1100             col[0].c   = 0;
1101             valA[0]    = -2.0 / (hx * hx) + -2.0 / (hy * hy) - 1.0 / (hz * hz);
1102             col[1].i   = ex;
1103             col[1].j   = ey - 1;
1104             col[1].k   = ez;
1105             col[1].loc = DOWN;
1106             col[1].c   = 0;
1107             valA[1]    = 1.0 / (hy * hy);
1108             col[2].i   = ex;
1109             col[2].j   = ey + 1;
1110             col[2].k   = ez;
1111             col[2].loc = DOWN;
1112             col[2].c   = 0;
1113             valA[2]    = 1.0 / (hy * hy);
1114             col[3].i   = ex - 1;
1115             col[3].j   = ey;
1116             col[3].k   = ez;
1117             col[3].loc = DOWN;
1118             col[3].c   = 0;
1119             valA[3]    = 1.0 / (hx * hx);
1120             col[4].i   = ex + 1;
1121             col[4].j   = ey;
1122             col[4].k   = ez;
1123             col[4].loc = DOWN;
1124             col[4].c   = 0;
1125             valA[4]    = 1.0 / (hx * hx);
1126             col[5].i   = ex;
1127             col[5].j   = ey;
1128             col[5].k   = ez - 1;
1129             col[5].loc = DOWN;
1130             col[5].c   = 0;
1131             valA[5]    = 1.0 / (hz * hz);
1132             /* Front term missing */
1133             col[6].i   = ex;
1134             col[6].j   = ey - 1;
1135             col[6].k   = ez;
1136             col[6].loc = ELEMENT;
1137             col[6].c   = 0;
1138             valA[6]    = 1.0 / hy;
1139             col[7].i   = ex;
1140             col[7].j   = ey;
1141             col[7].k   = ez;
1142             col[7].loc = ELEMENT;
1143             col[7].c   = 0;
1144             valA[7]    = -1.0 / hy;
1145           } else {
1146             nEntries   = 9;
1147             col[0].i   = ex;
1148             col[0].j   = ey;
1149             col[0].k   = ez;
1150             col[0].loc = DOWN;
1151             col[0].c   = 0;
1152             valA[0]    = -2.0 / (hx * hx) + -2.0 / (hy * hy) - 2.0 / (hz * hz);
1153             col[1].i   = ex;
1154             col[1].j   = ey - 1;
1155             col[1].k   = ez;
1156             col[1].loc = DOWN;
1157             col[1].c   = 0;
1158             valA[1]    = 1.0 / (hy * hy);
1159             col[2].i   = ex;
1160             col[2].j   = ey + 1;
1161             col[2].k   = ez;
1162             col[2].loc = DOWN;
1163             col[2].c   = 0;
1164             valA[2]    = 1.0 / (hy * hy);
1165             col[3].i   = ex - 1;
1166             col[3].j   = ey;
1167             col[3].k   = ez;
1168             col[3].loc = DOWN;
1169             col[3].c   = 0;
1170             valA[3]    = 1.0 / (hx * hx);
1171             col[4].i   = ex + 1;
1172             col[4].j   = ey;
1173             col[4].k   = ez;
1174             col[4].loc = DOWN;
1175             col[4].c   = 0;
1176             valA[4]    = 1.0 / (hx * hx);
1177             col[5].i   = ex;
1178             col[5].j   = ey;
1179             col[5].k   = ez - 1;
1180             col[5].loc = DOWN;
1181             col[5].c   = 0;
1182             valA[5]    = 1.0 / (hz * hz);
1183             col[6].i   = ex;
1184             col[6].j   = ey;
1185             col[6].k   = ez + 1;
1186             col[6].loc = DOWN;
1187             col[6].c   = 0;
1188             valA[6]    = 1.0 / (hz * hz);
1189             col[7].i   = ex;
1190             col[7].j   = ey - 1;
1191             col[7].k   = ez;
1192             col[7].loc = ELEMENT;
1193             col[7].c   = 0;
1194             valA[7]    = 1.0 / hy;
1195             col[8].i   = ex;
1196             col[8].j   = ey;
1197             col[8].k   = ez;
1198             col[8].loc = ELEMENT;
1199             col[8].c   = 0;
1200             valA[8]    = -1.0 / hy;
1201           }
1202           PetscCall(DMStagMatSetValuesStencil(dmSol, A, 1, &row, nEntries, col, valA, INSERT_VALUES));
1203           valRhs = fy(arrCoord[ez][ey][ex][icuy[0]], arrCoord[ez][ey][ex][icuy[1]], arrCoord[ez][ey][ex][icuy[2]]);
1204           PetscCall(DMStagVecSetValuesStencil(dmSol, rhs, 1, &row, &valRhs, INSERT_VALUES));
1205         }
1206 
1207         /* Equation on back face of this element */
1208         if (ez == 0) {
1209           /* Back boundary velocity Dirichlet */
1210           DMStagStencil     row;
1211           PetscScalar       valRhs;
1212           const PetscScalar valA = 1.0;
1213           row.i                  = ex;
1214           row.j                  = ey;
1215           row.k                  = ez;
1216           row.loc                = BACK;
1217           row.c                  = 0;
1218           PetscCall(DMStagMatSetValuesStencil(dmSol, A, 1, &row, 1, &row, &valA, INSERT_VALUES));
1219           valRhs = uzRef(arrCoord[ez][ey][ex][icuz[0]], arrCoord[ez][ey][ex][icuz[1]], arrCoord[ez][ey][ex][icuz[2]]);
1220           PetscCall(DMStagVecSetValuesStencil(dmSol, rhs, 1, &row, &valRhs, INSERT_VALUES));
1221         } else {
1222           /* Z-momentum equation, (w_xx + w_yy + w_zz) - p_z = f^z */
1223           DMStagStencil row, col[9];
1224           PetscScalar   valA[9], valRhs;
1225           PetscInt      nEntries;
1226 
1227           row.i   = ex;
1228           row.j   = ey;
1229           row.k   = ez;
1230           row.loc = BACK;
1231           row.c   = 0;
1232           if (ex == 0) {
1233             if (ey == 0) {
1234               nEntries   = 7;
1235               col[0].i   = ex;
1236               col[0].j   = ey;
1237               col[0].k   = ez;
1238               col[0].loc = BACK;
1239               col[0].c   = 0;
1240               valA[0]    = -1.0 / (hx * hx) + -1.0 / (hy * hy) - 2.0 / (hz * hz);
1241               /* Down term missing */
1242               col[1].i   = ex;
1243               col[1].j   = ey + 1;
1244               col[1].k   = ez;
1245               col[1].loc = BACK;
1246               col[1].c   = 0;
1247               valA[1]    = 1.0 / (hy * hy);
1248               /* Left term missing */
1249               col[2].i   = ex + 1;
1250               col[2].j   = ey;
1251               col[2].k   = ez;
1252               col[2].loc = BACK;
1253               col[2].c   = 0;
1254               valA[2]    = 1.0 / (hx * hx);
1255               col[3].i   = ex;
1256               col[3].j   = ey;
1257               col[3].k   = ez - 1;
1258               col[3].loc = BACK;
1259               col[3].c   = 0;
1260               valA[3]    = 1.0 / (hz * hz);
1261               col[4].i   = ex;
1262               col[4].j   = ey;
1263               col[4].k   = ez + 1;
1264               col[4].loc = BACK;
1265               col[4].c   = 0;
1266               valA[4]    = 1.0 / (hz * hz);
1267               col[5].i   = ex;
1268               col[5].j   = ey;
1269               col[5].k   = ez - 1;
1270               col[5].loc = ELEMENT;
1271               col[5].c   = 0;
1272               valA[5]    = 1.0 / hz;
1273               col[6].i   = ex;
1274               col[6].j   = ey;
1275               col[6].k   = ez;
1276               col[6].loc = ELEMENT;
1277               col[6].c   = 0;
1278               valA[6]    = -1.0 / hz;
1279             } else if (ey == N[1] - 1) {
1280               nEntries   = 7;
1281               col[0].i   = ex;
1282               col[0].j   = ey;
1283               col[0].k   = ez;
1284               col[0].loc = BACK;
1285               col[0].c   = 0;
1286               valA[0]    = -1.0 / (hx * hx) + -1.0 / (hy * hy) - 2.0 / (hz * hz);
1287               col[1].i   = ex;
1288               col[1].j   = ey - 1;
1289               col[1].k   = ez;
1290               col[1].loc = BACK;
1291               col[1].c   = 0;
1292               valA[1]    = 1.0 / (hy * hy);
1293               /* Up term missing */
1294               /* Left term missing */
1295               col[2].i   = ex + 1;
1296               col[2].j   = ey;
1297               col[2].k   = ez;
1298               col[2].loc = BACK;
1299               col[2].c   = 0;
1300               valA[2]    = 1.0 / (hx * hx);
1301               col[3].i   = ex;
1302               col[3].j   = ey;
1303               col[3].k   = ez - 1;
1304               col[3].loc = BACK;
1305               col[3].c   = 0;
1306               valA[3]    = 1.0 / (hz * hz);
1307               col[4].i   = ex;
1308               col[4].j   = ey;
1309               col[4].k   = ez + 1;
1310               col[4].loc = BACK;
1311               col[4].c   = 0;
1312               valA[4]    = 1.0 / (hz * hz);
1313               col[5].i   = ex;
1314               col[5].j   = ey;
1315               col[5].k   = ez - 1;
1316               col[5].loc = ELEMENT;
1317               col[5].c   = 0;
1318               valA[5]    = 1.0 / hz;
1319               col[6].i   = ex;
1320               col[6].j   = ey;
1321               col[6].k   = ez;
1322               col[6].loc = ELEMENT;
1323               col[6].c   = 0;
1324               valA[6]    = -1.0 / hz;
1325             } else {
1326               nEntries   = 8;
1327               col[0].i   = ex;
1328               col[0].j   = ey;
1329               col[0].k   = ez;
1330               col[0].loc = BACK;
1331               col[0].c   = 0;
1332               valA[0]    = -1.0 / (hx * hx) + -2.0 / (hy * hy) - 2.0 / (hz * hz);
1333               col[1].i   = ex;
1334               col[1].j   = ey - 1;
1335               col[1].k   = ez;
1336               col[1].loc = BACK;
1337               col[1].c   = 0;
1338               valA[1]    = 1.0 / (hy * hy);
1339               col[2].i   = ex;
1340               col[2].j   = ey + 1;
1341               col[2].k   = ez;
1342               col[2].loc = BACK;
1343               col[2].c   = 0;
1344               valA[2]    = 1.0 / (hy * hy);
1345               /* Left term missing */
1346               col[3].i   = ex + 1;
1347               col[3].j   = ey;
1348               col[3].k   = ez;
1349               col[3].loc = BACK;
1350               col[3].c   = 0;
1351               valA[3]    = 1.0 / (hx * hx);
1352               col[4].i   = ex;
1353               col[4].j   = ey;
1354               col[4].k   = ez - 1;
1355               col[4].loc = BACK;
1356               col[4].c   = 0;
1357               valA[4]    = 1.0 / (hz * hz);
1358               col[5].i   = ex;
1359               col[5].j   = ey;
1360               col[5].k   = ez + 1;
1361               col[5].loc = BACK;
1362               col[5].c   = 0;
1363               valA[5]    = 1.0 / (hz * hz);
1364               col[6].i   = ex;
1365               col[6].j   = ey;
1366               col[6].k   = ez - 1;
1367               col[6].loc = ELEMENT;
1368               col[6].c   = 0;
1369               valA[6]    = 1.0 / hz;
1370               col[7].i   = ex;
1371               col[7].j   = ey;
1372               col[7].k   = ez;
1373               col[7].loc = ELEMENT;
1374               col[7].c   = 0;
1375               valA[7]    = -1.0 / hz;
1376             }
1377           } else if (ex == N[0] - 1) {
1378             if (ey == 0) {
1379               nEntries   = 7;
1380               col[0].i   = ex;
1381               col[0].j   = ey;
1382               col[0].k   = ez;
1383               col[0].loc = BACK;
1384               col[0].c   = 0;
1385               valA[0]    = -1.0 / (hx * hx) + -1.0 / (hy * hy) - 2.0 / (hz * hz);
1386               /* Down term missing */
1387               col[1].i   = ex;
1388               col[1].j   = ey + 1;
1389               col[1].k   = ez;
1390               col[1].loc = BACK;
1391               col[1].c   = 0;
1392               valA[1]    = 1.0 / (hy * hy);
1393               col[2].i   = ex - 1;
1394               col[2].j   = ey;
1395               col[2].k   = ez;
1396               col[2].loc = BACK;
1397               col[2].c   = 0;
1398               valA[2]    = 1.0 / (hx * hx);
1399               /* Right term missing */
1400               col[3].i   = ex;
1401               col[3].j   = ey;
1402               col[3].k   = ez - 1;
1403               col[3].loc = BACK;
1404               col[3].c   = 0;
1405               valA[3]    = 1.0 / (hz * hz);
1406               col[4].i   = ex;
1407               col[4].j   = ey;
1408               col[4].k   = ez + 1;
1409               col[4].loc = BACK;
1410               col[4].c   = 0;
1411               valA[4]    = 1.0 / (hz * hz);
1412               col[5].i   = ex;
1413               col[5].j   = ey;
1414               col[5].k   = ez - 1;
1415               col[5].loc = ELEMENT;
1416               col[5].c   = 0;
1417               valA[5]    = 1.0 / hz;
1418               col[6].i   = ex;
1419               col[6].j   = ey;
1420               col[6].k   = ez;
1421               col[6].loc = ELEMENT;
1422               col[6].c   = 0;
1423               valA[6]    = -1.0 / hz;
1424             } else if (ey == N[1] - 1) {
1425               nEntries   = 7;
1426               col[0].i   = ex;
1427               col[0].j   = ey;
1428               col[0].k   = ez;
1429               col[0].loc = BACK;
1430               col[0].c   = 0;
1431               valA[0]    = -1.0 / (hx * hx) + -1.0 / (hy * hy) - 2.0 / (hz * hz);
1432               col[1].i   = ex;
1433               col[1].j   = ey - 1;
1434               col[1].k   = ez;
1435               col[1].loc = BACK;
1436               col[1].c   = 0;
1437               valA[1]    = 1.0 / (hy * hy);
1438               /* Up term missing */
1439               col[2].i   = ex - 1;
1440               col[2].j   = ey;
1441               col[2].k   = ez;
1442               col[2].loc = BACK;
1443               col[2].c   = 0;
1444               valA[2]    = 1.0 / (hx * hx);
1445               /* Right term missing */
1446               col[3].i   = ex;
1447               col[3].j   = ey;
1448               col[3].k   = ez - 1;
1449               col[3].loc = BACK;
1450               col[3].c   = 0;
1451               valA[3]    = 1.0 / (hz * hz);
1452               col[4].i   = ex;
1453               col[4].j   = ey;
1454               col[4].k   = ez + 1;
1455               col[4].loc = BACK;
1456               col[4].c   = 0;
1457               valA[4]    = 1.0 / (hz * hz);
1458               col[5].i   = ex;
1459               col[5].j   = ey;
1460               col[5].k   = ez - 1;
1461               col[5].loc = ELEMENT;
1462               col[5].c   = 0;
1463               valA[5]    = 1.0 / hz;
1464               col[6].i   = ex;
1465               col[6].j   = ey;
1466               col[6].k   = ez;
1467               col[6].loc = ELEMENT;
1468               col[6].c   = 0;
1469               valA[6]    = -1.0 / hz;
1470             } else {
1471               nEntries   = 8;
1472               col[0].i   = ex;
1473               col[0].j   = ey;
1474               col[0].k   = ez;
1475               col[0].loc = BACK;
1476               col[0].c   = 0;
1477               valA[0]    = -1.0 / (hx * hx) + -2.0 / (hy * hy) - 2.0 / (hz * hz);
1478               col[1].i   = ex;
1479               col[1].j   = ey - 1;
1480               col[1].k   = ez;
1481               col[1].loc = BACK;
1482               col[1].c   = 0;
1483               valA[1]    = 1.0 / (hy * hy);
1484               col[2].i   = ex;
1485               col[2].j   = ey + 1;
1486               col[2].k   = ez;
1487               col[2].loc = BACK;
1488               col[2].c   = 0;
1489               valA[2]    = 1.0 / (hy * hy);
1490               col[3].i   = ex - 1;
1491               col[3].j   = ey;
1492               col[3].k   = ez;
1493               col[3].loc = BACK;
1494               col[3].c   = 0;
1495               valA[3]    = 1.0 / (hx * hx);
1496               /* Right term missing */
1497               col[4].i   = ex;
1498               col[4].j   = ey;
1499               col[4].k   = ez - 1;
1500               col[4].loc = BACK;
1501               col[4].c   = 0;
1502               valA[4]    = 1.0 / (hz * hz);
1503               col[5].i   = ex;
1504               col[5].j   = ey;
1505               col[5].k   = ez + 1;
1506               col[5].loc = BACK;
1507               col[5].c   = 0;
1508               valA[5]    = 1.0 / (hz * hz);
1509               col[6].i   = ex;
1510               col[6].j   = ey;
1511               col[6].k   = ez - 1;
1512               col[6].loc = ELEMENT;
1513               col[6].c   = 0;
1514               valA[6]    = 1.0 / hz;
1515               col[7].i   = ex;
1516               col[7].j   = ey;
1517               col[7].k   = ez;
1518               col[7].loc = ELEMENT;
1519               col[7].c   = 0;
1520               valA[7]    = -1.0 / hz;
1521             }
1522           } else if (ey == 0) {
1523             nEntries   = 8;
1524             col[0].i   = ex;
1525             col[0].j   = ey;
1526             col[0].k   = ez;
1527             col[0].loc = BACK;
1528             col[0].c   = 0;
1529             valA[0]    = -2.0 / (hx * hx) + -1.0 / (hy * hy) - 2.0 / (hz * hz);
1530             /* Down term missing */
1531             col[1].i   = ex;
1532             col[1].j   = ey + 1;
1533             col[1].k   = ez;
1534             col[1].loc = BACK;
1535             col[1].c   = 0;
1536             valA[1]    = 1.0 / (hy * hy);
1537             col[2].i   = ex - 1;
1538             col[2].j   = ey;
1539             col[2].k   = ez;
1540             col[2].loc = BACK;
1541             col[2].c   = 0;
1542             valA[2]    = 1.0 / (hx * hx);
1543             col[3].i   = ex + 1;
1544             col[3].j   = ey;
1545             col[3].k   = ez;
1546             col[3].loc = BACK;
1547             col[3].c   = 0;
1548             valA[3]    = 1.0 / (hx * hx);
1549             col[4].i   = ex;
1550             col[4].j   = ey;
1551             col[4].k   = ez - 1;
1552             col[4].loc = BACK;
1553             col[4].c   = 0;
1554             valA[4]    = 1.0 / (hz * hz);
1555             col[5].i   = ex;
1556             col[5].j   = ey;
1557             col[5].k   = ez + 1;
1558             col[5].loc = BACK;
1559             col[5].c   = 0;
1560             valA[5]    = 1.0 / (hz * hz);
1561             col[6].i   = ex;
1562             col[6].j   = ey;
1563             col[6].k   = ez - 1;
1564             col[6].loc = ELEMENT;
1565             col[6].c   = 0;
1566             valA[6]    = 1.0 / hz;
1567             col[7].i   = ex;
1568             col[7].j   = ey;
1569             col[7].k   = ez;
1570             col[7].loc = ELEMENT;
1571             col[7].c   = 0;
1572             valA[7]    = -1.0 / hz;
1573           } else if (ey == N[1] - 1) {
1574             nEntries   = 8;
1575             col[0].i   = ex;
1576             col[0].j   = ey;
1577             col[0].k   = ez;
1578             col[0].loc = BACK;
1579             col[0].c   = 0;
1580             valA[0]    = -2.0 / (hx * hx) + -1.0 / (hy * hy) - 2.0 / (hz * hz);
1581             col[1].i   = ex;
1582             col[1].j   = ey - 1;
1583             col[1].k   = ez;
1584             col[1].loc = BACK;
1585             col[1].c   = 0;
1586             valA[1]    = 1.0 / (hy * hy);
1587             /* Up term missing */
1588             col[2].i   = ex - 1;
1589             col[2].j   = ey;
1590             col[2].k   = ez;
1591             col[2].loc = BACK;
1592             col[2].c   = 0;
1593             valA[2]    = 1.0 / (hx * hx);
1594             col[3].i   = ex + 1;
1595             col[3].j   = ey;
1596             col[3].k   = ez;
1597             col[3].loc = BACK;
1598             col[3].c   = 0;
1599             valA[3]    = 1.0 / (hx * hx);
1600             col[4].i   = ex;
1601             col[4].j   = ey;
1602             col[4].k   = ez - 1;
1603             col[4].loc = BACK;
1604             col[4].c   = 0;
1605             valA[4]    = 1.0 / (hz * hz);
1606             col[5].i   = ex;
1607             col[5].j   = ey;
1608             col[5].k   = ez + 1;
1609             col[5].loc = BACK;
1610             col[5].c   = 0;
1611             valA[5]    = 1.0 / (hz * hz);
1612             col[6].i   = ex;
1613             col[6].j   = ey;
1614             col[6].k   = ez - 1;
1615             col[6].loc = ELEMENT;
1616             col[6].c   = 0;
1617             valA[6]    = 1.0 / hz;
1618             col[7].i   = ex;
1619             col[7].j   = ey;
1620             col[7].k   = ez;
1621             col[7].loc = ELEMENT;
1622             col[7].c   = 0;
1623             valA[7]    = -1.0 / hz;
1624           } else {
1625             nEntries   = 9;
1626             col[0].i   = ex;
1627             col[0].j   = ey;
1628             col[0].k   = ez;
1629             col[0].loc = BACK;
1630             col[0].c   = 0;
1631             valA[0]    = -2.0 / (hx * hx) + -2.0 / (hy * hy) - 2.0 / (hz * hz);
1632             col[1].i   = ex;
1633             col[1].j   = ey - 1;
1634             col[1].k   = ez;
1635             col[1].loc = BACK;
1636             col[1].c   = 0;
1637             valA[1]    = 1.0 / (hy * hy);
1638             col[2].i   = ex;
1639             col[2].j   = ey + 1;
1640             col[2].k   = ez;
1641             col[2].loc = BACK;
1642             col[2].c   = 0;
1643             valA[2]    = 1.0 / (hy * hy);
1644             col[3].i   = ex - 1;
1645             col[3].j   = ey;
1646             col[3].k   = ez;
1647             col[3].loc = BACK;
1648             col[3].c   = 0;
1649             valA[3]    = 1.0 / (hx * hx);
1650             col[4].i   = ex + 1;
1651             col[4].j   = ey;
1652             col[4].k   = ez;
1653             col[4].loc = BACK;
1654             col[4].c   = 0;
1655             valA[4]    = 1.0 / (hx * hx);
1656             col[5].i   = ex;
1657             col[5].j   = ey;
1658             col[5].k   = ez - 1;
1659             col[5].loc = BACK;
1660             col[5].c   = 0;
1661             valA[5]    = 1.0 / (hz * hz);
1662             col[6].i   = ex;
1663             col[6].j   = ey;
1664             col[6].k   = ez + 1;
1665             col[6].loc = BACK;
1666             col[6].c   = 0;
1667             valA[6]    = 1.0 / (hz * hz);
1668             col[7].i   = ex;
1669             col[7].j   = ey;
1670             col[7].k   = ez - 1;
1671             col[7].loc = ELEMENT;
1672             col[7].c   = 0;
1673             valA[7]    = 1.0 / hz;
1674             col[8].i   = ex;
1675             col[8].j   = ey;
1676             col[8].k   = ez;
1677             col[8].loc = ELEMENT;
1678             col[8].c   = 0;
1679             valA[8]    = -1.0 / hz;
1680           }
1681           PetscCall(DMStagMatSetValuesStencil(dmSol, A, 1, &row, nEntries, col, valA, INSERT_VALUES));
1682           valRhs = fz(arrCoord[ez][ey][ex][icuz[0]], arrCoord[ez][ey][ex][icuz[1]], arrCoord[ez][ey][ex][icuz[2]]);
1683           PetscCall(DMStagVecSetValuesStencil(dmSol, rhs, 1, &row, &valRhs, INSERT_VALUES));
1684         }
1685 
1686         /* P equation : u_x + v_y + w_z = g
1687            Note that this includes an explicit zero on the diagonal. This is only needed for
1688            direct solvers (not required if using an iterative solver and setting the constant-pressure nullspace) */
1689         if (pinPressure && ex == 0 && ey == 0 && ez == 0) { /* Pin the first pressure node, if requested */
1690           DMStagStencil row;
1691           PetscScalar   valA, valRhs;
1692           row.i   = ex;
1693           row.j   = ey;
1694           row.k   = ez;
1695           row.loc = ELEMENT;
1696           row.c   = 0;
1697           valA    = 1.0;
1698           PetscCall(DMStagMatSetValuesStencil(dmSol, A, 1, &row, 1, &row, &valA, INSERT_VALUES));
1699           valRhs = pRef(arrCoord[ez][ey][ex][icp[0]], arrCoord[ez][ey][ex][icp[1]], arrCoord[ez][ey][ex][icp[2]]);
1700           PetscCall(DMStagVecSetValuesStencil(dmSol, rhs, 1, &row, &valRhs, INSERT_VALUES));
1701         } else {
1702           DMStagStencil row, col[7];
1703           PetscScalar   valA[7], valRhs;
1704 
1705           row.i      = ex;
1706           row.j      = ey;
1707           row.k      = ez;
1708           row.loc    = ELEMENT;
1709           row.c      = 0;
1710           col[0].i   = ex;
1711           col[0].j   = ey;
1712           col[0].k   = ez;
1713           col[0].loc = LEFT;
1714           col[0].c   = 0;
1715           valA[0]    = -1.0 / hx;
1716           col[1].i   = ex;
1717           col[1].j   = ey;
1718           col[1].k   = ez;
1719           col[1].loc = RIGHT;
1720           col[1].c   = 0;
1721           valA[1]    = 1.0 / hx;
1722           col[2].i   = ex;
1723           col[2].j   = ey;
1724           col[2].k   = ez;
1725           col[2].loc = DOWN;
1726           col[2].c   = 0;
1727           valA[2]    = -1.0 / hy;
1728           col[3].i   = ex;
1729           col[3].j   = ey;
1730           col[3].k   = ez;
1731           col[3].loc = UP;
1732           col[3].c   = 0;
1733           valA[3]    = 1.0 / hy;
1734           col[4].i   = ex;
1735           col[4].j   = ey;
1736           col[4].k   = ez;
1737           col[4].loc = BACK;
1738           col[4].c   = 0;
1739           valA[4]    = -1.0 / hz;
1740           col[5].i   = ex;
1741           col[5].j   = ey;
1742           col[5].k   = ez;
1743           col[5].loc = FRONT;
1744           col[5].c   = 0;
1745           valA[5]    = 1.0 / hz;
1746           col[6]     = row;
1747           valA[6]    = 0.0;
1748           PetscCall(DMStagMatSetValuesStencil(dmSol, A, 1, &row, 7, col, valA, INSERT_VALUES));
1749           valRhs = g(arrCoord[ez][ey][ex][icp[0]], arrCoord[ez][ey][ex][icp[1]], arrCoord[ez][ey][ex][icp[2]]);
1750           PetscCall(DMStagVecSetValuesStencil(dmSol, rhs, 1, &row, &valRhs, INSERT_VALUES));
1751         }
1752       }
1753     }
1754   }
1755   PetscCall(DMStagVecRestoreArrayRead(dmCoord, coordLocal, &arrCoord));
1756   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
1757   PetscCall(VecAssemblyBegin(rhs));
1758   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
1759   PetscCall(VecAssemblyEnd(rhs));
1760   PetscFunctionReturn(PETSC_SUCCESS);
1761 }
1762 
1763 /* Create a pressure-only DMStag and use it to generate a nullspace vector
1764    - Create a compatible DMStag with one dof per element (and nothing else).
1765    - Create a constant vector and normalize it
1766    - Migrate it to a vector on the original dmSol (making use of the fact
1767    that this will fill in zeros for "extra" dof)
1768    - Set the nullspace for the operator
1769    - Destroy everything (the operator keeps the references it needs) */
AttachNullspace(DM dmSol,Mat A)1770 static PetscErrorCode AttachNullspace(DM dmSol, Mat A)
1771 {
1772   DM           dmPressure;
1773   Vec          constantPressure, basis;
1774   PetscReal    nrm;
1775   MatNullSpace matNullSpace;
1776 
1777   PetscFunctionBeginUser;
1778   PetscCall(DMStagCreateCompatibleDMStag(dmSol, 0, 0, 0, 1, &dmPressure));
1779   PetscCall(DMGetGlobalVector(dmPressure, &constantPressure));
1780   PetscCall(VecSet(constantPressure, 1.0));
1781   PetscCall(VecNorm(constantPressure, NORM_2, &nrm));
1782   PetscCall(VecScale(constantPressure, 1.0 / nrm));
1783   PetscCall(DMCreateGlobalVector(dmSol, &basis));
1784   PetscCall(DMStagMigrateVec(dmPressure, constantPressure, dmSol, basis));
1785   PetscCall(MatNullSpaceCreate(PetscObjectComm((PetscObject)dmSol), PETSC_FALSE, 1, &basis, &matNullSpace));
1786   PetscCall(VecDestroy(&basis));
1787   PetscCall(VecDestroy(&constantPressure));
1788   PetscCall(MatSetNullSpace(A, matNullSpace));
1789   PetscCall(MatNullSpaceDestroy(&matNullSpace));
1790   PetscFunctionReturn(PETSC_SUCCESS);
1791 }
1792 
CreateReferenceSolution(DM dmSol,Vec * pSolRef)1793 static PetscErrorCode CreateReferenceSolution(DM dmSol, Vec *pSolRef)
1794 {
1795   PetscInt        start[3], n[3], nExtra[3], ex, ey, ez, d;
1796   PetscInt        ip, iux, iuy, iuz, icp[3], icux[3], icuy[3], icuz[3];
1797   Vec             solRef, solRefLocal, coord, coordLocal;
1798   DM              dmCoord;
1799   PetscScalar ****arrSol, ****arrCoord;
1800 
1801   PetscFunctionBeginUser;
1802   PetscCall(DMCreateGlobalVector(dmSol, pSolRef));
1803   solRef = *pSolRef;
1804   PetscCall(DMStagGetCorners(dmSol, &start[0], &start[1], &start[2], &n[0], &n[1], &n[2], &nExtra[0], &nExtra[1], &nExtra[2]));
1805   PetscCall(DMGetCoordinateDM(dmSol, &dmCoord));
1806   PetscCall(DMGetCoordinates(dmSol, &coord));
1807   PetscCall(DMGetLocalVector(dmCoord, &coordLocal));
1808   PetscCall(DMGlobalToLocal(dmCoord, coord, INSERT_VALUES, coordLocal));
1809   PetscCall(DMStagGetLocationSlot(dmSol, ELEMENT, 0, &ip));
1810   PetscCall(DMStagGetLocationSlot(dmSol, LEFT, 0, &iux));
1811   PetscCall(DMStagGetLocationSlot(dmSol, DOWN, 0, &iuy));
1812   PetscCall(DMStagGetLocationSlot(dmSol, BACK, 0, &iuz));
1813   for (d = 0; d < 3; ++d) {
1814     PetscCall(DMStagGetLocationSlot(dmCoord, ELEMENT, d, &icp[d]));
1815     PetscCall(DMStagGetLocationSlot(dmCoord, LEFT, d, &icux[d]));
1816     PetscCall(DMStagGetLocationSlot(dmCoord, DOWN, d, &icuy[d]));
1817     PetscCall(DMStagGetLocationSlot(dmCoord, BACK, d, &icuz[d]));
1818   }
1819   PetscCall(DMStagVecGetArrayRead(dmCoord, coordLocal, &arrCoord));
1820   PetscCall(DMGetLocalVector(dmSol, &solRefLocal));
1821   PetscCall(DMStagVecGetArray(dmSol, solRefLocal, &arrSol));
1822   for (ez = start[2]; ez < start[2] + n[2] + nExtra[2]; ++ez) {
1823     for (ey = start[1]; ey < start[1] + n[1] + nExtra[1]; ++ey) {
1824       for (ex = start[0]; ex < start[0] + n[0] + nExtra[0]; ++ex) {
1825         if (ex < start[0] + n[0] && ey < start[1] + n[1]) arrSol[ez][ey][ex][iuz] = uzRef(arrCoord[ez][ey][ex][icuz[0]], arrCoord[ez][ey][ex][icuz[1]], arrCoord[ez][ey][ex][icuz[2]]);
1826         if (ex < start[0] + n[0] && ey < start[2] + n[2]) arrSol[ez][ey][ex][iuy] = uyRef(arrCoord[ez][ey][ex][icuy[0]], arrCoord[ez][ey][ex][icuy[1]], arrCoord[ez][ey][ex][icuy[2]]);
1827         if (ex < start[1] + n[1] && ey < start[2] + n[2]) arrSol[ez][ey][ex][iux] = uxRef(arrCoord[ez][ey][ex][icux[0]], arrCoord[ez][ey][ex][icux[1]], arrCoord[ez][ey][ex][icux[2]]);
1828         if (ex < start[0] + n[0] && ey < start[1] + n[1] && ez < start[2] + n[2]) arrSol[ez][ey][ex][ip] = pRef(arrCoord[ez][ey][ex][icp[0]], arrCoord[ez][ey][ex][icp[1]], arrCoord[ez][ey][ex][icp[2]]);
1829       }
1830     }
1831   }
1832   PetscCall(DMStagVecRestoreArrayRead(dmCoord, coordLocal, &arrCoord));
1833   PetscCall(DMStagVecRestoreArray(dmSol, solRefLocal, &arrSol));
1834   PetscCall(DMLocalToGlobal(dmSol, solRefLocal, INSERT_VALUES, solRef));
1835   PetscCall(DMRestoreLocalVector(dmCoord, &coordLocal));
1836   PetscCall(DMRestoreLocalVector(dmSol, &solRefLocal));
1837   PetscFunctionReturn(PETSC_SUCCESS);
1838 }
1839 
CheckSolution(Vec sol,Vec solRef)1840 static PetscErrorCode CheckSolution(Vec sol, Vec solRef)
1841 {
1842   Vec       diff;
1843   PetscReal normsolRef, errAbs, errRel;
1844 
1845   PetscFunctionBeginUser;
1846   PetscCall(VecDuplicate(sol, &diff));
1847   PetscCall(VecCopy(sol, diff));
1848   PetscCall(VecAXPY(diff, -1.0, solRef));
1849   PetscCall(VecNorm(diff, NORM_2, &errAbs));
1850   PetscCall(VecNorm(solRef, NORM_2, &normsolRef));
1851   errRel = errAbs / normsolRef;
1852   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Error (abs): %g\nError (rel): %g\n", (double)errAbs, (double)errRel));
1853   PetscCall(VecDestroy(&diff));
1854   PetscFunctionReturn(PETSC_SUCCESS);
1855 }
1856 
1857 /*TEST
1858 
1859    test:
1860       suffix: 1
1861       requires: mumps
1862       nsize: 27
1863       args: -ksp_monitor_short -ksp_converged_reason -stag_ranks_x 3 -stag_ranks_y 3 -stag_ranks_z 3 -pc_fieldsplit_schur_fact_type diag -fieldsplit_0_ksp_type preonly -fieldsplit_1_pc_type none -fieldsplit_0_pc_type lu -fieldsplit_1_ksp_type gmres -fieldsplit_1_ksp_max_it 20
1864 
1865    test:
1866       suffix: 2
1867       requires: !single
1868       nsize: 4
1869       args: -ksp_monitor_short -ksp_converged_reason -pc_fieldsplit_schur_fact_type diag -fieldsplit_0_ksp_type preonly -fieldsplit_1_pc_type none -fieldsplit_0_pc_type gamg -fieldsplit_0_ksp_type gmres -fieldsplit_0_pc_gamg_esteig_ksp_max_it 5
1870 
1871    test:
1872       suffix: direct_umfpack
1873       requires: suitesparse
1874       nsize: 1
1875       args: -pinpressure 1 -stag_grid_x 5 -stag_grid_y 3 -stag_grid_z 4 -ksp_monitor_short -pc_type lu -pc_factor_mat_solver_type umfpack
1876 
1877 TEST*/
1878