Actual source code: ex1.c

```  1: static char help[] = "Solve a toy 1D problem on a staggered grid.\n\
2:                       Accepts command line options -a, -b, and -c\n\
3:                       and approximately solves\n\
4:                       u''(x) = c, u(0) = 1, u(1) = b\n\n";
5: /*

7:    To demonstrate the basic functionality of DMStag, solves a second-order ODE,

9:        u''(x) = f(x),  0 < x < 1
10:        u(0) = a
11:        u(1) = b

13:    in mixed form, that is by introducing an auxiliary variable p

15:       p'(x) = f(x), p - u'(x) = 0, 0 < x < 1
16:       u(0) = a
17:       u(1) = b

19:    For f == c, the solution is

21:      u(x) = a + (b - a - (c/2)) x + (c/2) x^2
22:      p(x) = (b - a - (c/2)) + c x

24:    To find an approximate solution, discretize by storing values of p in
25:    elements and values of u on their boundaries, and using first-order finite
26:    differences.

28:    This should in fact produce a (nodal) solution with no discretization error,
29:    so differences from the reference solution will be restricted to those induced
30:    by floating point operations (in particular, the finite differences) and the
31:    accuracy of the linear solve.

33:    Parameters for the main grid can be controlled with command line options, e.g.

35:      -stag_grid_x 10

37:   In particular to notice in this example are the two methods of indexing. The
38:   first is analogous to the use of MatStencil with DMDA, and the second is
39:   analogous to the use of DMDAVecGetArrayDOF().

41:   The first, recommended for ease of use, is based on naming an element in the
42:   global grid, a location in its support, and a component. For example,
43:   one might consider element e, the left side (a vertex in 1d), and the first
44:   component (index 0). This is accomplished by populating a DMStagStencil struct,
45:   e.g.

47:       DMStagStencil stencil;
48:       stencil.i   = i
49:       stencil.loc = DMSTAG_LEFT;
50:       stencil.c   = 0

52:   Note that below, for convenenience, we #define an alias LEFT for DMSTAG_LEFT.

54:   The second, which ensures maximum efficiency, makes use of the underlying
55:   block structure of local DMStag-derived vectors, and requires the user to
56:   obtain the correct local offset for the degrees of freedom they would like to
57:   use. This is made simple with the helper function DMStagGetLocationSlot().

59:   Note that the linear system being solved is indefinite, so is not entirely
60:   trivial to invert. The default solver here (GMRES/Jacobi) is a poor choice,
61:   made to avoid depending on an external package. To solve a larger system,
62:   the usual method for a 1-d problem such as this is to choose a sophisticated

65:     \$PETSC_DIR/\$PETSC_ARCH/bin/mpiexec -n 3 ./stag_ex2 -stag_grid_x 100 -pc_type lu -pc_factor_mat_solver_package umfpack

67:   You can also impose a periodic boundary condition, in which case -b and -c are
68:   ignored; b = a and c = 0.0 are used, giving a constant u == a , p == 0.

70:       -stag_boundary_type_x periodic

72: */
73: #include <petscdm.h>
74: #include <petscksp.h>
75: #include <petscdmstag.h>

77: /* Shorter, more convenient names for DMStagStencilLocation entries */
78: #define LEFT    DMSTAG_LEFT
79: #define RIGHT   DMSTAG_RIGHT
80: #define ELEMENT DMSTAG_ELEMENT

82: int main(int argc, char **argv)
83: {
84:   DM             dmSol, dmForcing;
85:   DM             dmCoordSol;
86:   Vec            sol, solRef, solRefLocal, f, fLocal, rhs, coordSolLocal;
87:   Mat            A;
88:   PetscScalar    a, b, c, h;
89:   KSP            ksp;
90:   PC             pc;
91:   PetscInt       start, n, e, nExtra;
92:   PetscInt       iu, ip, ixu, ixp;
93:   PetscBool      isLastRank, isFirstRank;
94:   PetscScalar  **arrSol, **arrCoordSol;
95:   DMBoundaryType boundary;

97:   const PetscReal domainSize = 1.0;

99:   /* Initialize PETSc */
100:   PetscFunctionBeginUser;
101:   PetscCall(PetscInitialize(&argc, &argv, (char *)0, help));

103:   /* Create 1D DMStag for the solution, and set up. Note that you can supply many
104:      command line options (see the man page for DMStagCreate1d)
105:   */
106:   PetscCall(DMStagCreate1d(PETSC_COMM_WORLD, DM_BOUNDARY_NONE, 3, 1, 1, DMSTAG_STENCIL_BOX, 1, NULL, &dmSol));
107:   PetscCall(DMSetFromOptions(dmSol));
108:   PetscCall(DMSetUp(dmSol));

110:   /* Create uniform coordinates. Note that in higher-dimensional examples,
111:       coordinates are created differently.*/
112:   PetscCall(DMStagSetUniformCoordinatesExplicit(dmSol, 0.0, domainSize, 0.0, 0.0, 0.0, 0.0));

114:   /* Determine boundary type */
115:   PetscCall(DMStagGetBoundaryTypes(dmSol, &boundary, NULL, NULL));

117:   /* Process command line options (depends on DMStag setup) */
118:   a = 1.0;
119:   b = 2.0;
120:   c = 1.0;
121:   PetscCall(PetscOptionsGetScalar(NULL, NULL, "-a", &a, NULL));
122:   if (boundary == DM_BOUNDARY_PERIODIC) {
123:     b = a;
124:     c = 0.0;
125:   } else {
126:     PetscCall(PetscOptionsGetScalar(NULL, NULL, "-b", &b, NULL));
127:     PetscCall(PetscOptionsGetScalar(NULL, NULL, "-c", &c, NULL));
128:   }

130:   /* Compute reference solution on the grid, using direct array access */
131:   PetscCall(DMCreateGlobalVector(dmSol, &solRef));
132:   PetscCall(DMGetLocalVector(dmSol, &solRefLocal));
133:   PetscCall(DMStagVecGetArray(dmSol, solRefLocal, &arrSol));
134:   PetscCall(DMGetCoordinateDM(dmSol, &dmCoordSol));
135:   PetscCall(DMGetCoordinatesLocal(dmSol, &coordSolLocal));
137:   PetscCall(DMStagGetCorners(dmSol, &start, NULL, NULL, &n, NULL, NULL, &nExtra, NULL, NULL));

139:   /* Get the correct entries for each of our variables in local element-wise storage */
140:   PetscCall(DMStagGetLocationSlot(dmSol, LEFT, 0, &iu));
141:   PetscCall(DMStagGetLocationSlot(dmSol, ELEMENT, 0, &ip));
142:   PetscCall(DMStagGetLocationSlot(dmCoordSol, LEFT, 0, &ixu));
143:   PetscCall(DMStagGetLocationSlot(dmCoordSol, ELEMENT, 0, &ixp));
144:   for (e = start; e < start + n + nExtra; ++e) {
145:     {
146:       const PetscScalar coordu = arrCoordSol[e][ixu];
147:       arrSol[e][iu]            = a + (b - a - (c / 2.0)) * coordu + (c / 2.0) * coordu * coordu;
148:     }
149:     if (e < start + n) {
150:       const PetscScalar coordp = arrCoordSol[e][ixp];
151:       arrSol[e][ip]            = b - a - (c / 2.0) + c * coordp;
152:     }
153:   }
155:   PetscCall(DMStagVecRestoreArray(dmSol, solRefLocal, &arrSol));
156:   PetscCall(DMLocalToGlobal(dmSol, solRefLocal, INSERT_VALUES, solRef));
157:   PetscCall(DMRestoreLocalVector(dmSol, &solRefLocal));

159:   /* Create another 1D DMStag for the forcing term, and populate a field on it.
160:      Here this is not really necessary, but in other contexts we may have auxiliary
161:      fields which we use to construct the linear system.

163:      This second DM represents the same physical domain, but has a different
164:      "default section" (though the current implementation does NOT actually use
165:      PetscSection). Since it is created as a derivative of the original DMStag,
166:      we can be confident that it is compatible. One could check with DMGetCompatiblity() */
167:   PetscCall(DMStagCreateCompatibleDMStag(dmSol, 1, 0, 0, 0, &dmForcing));
168:   PetscCall(DMCreateGlobalVector(dmForcing, &f));
169:   PetscCall(VecSet(f, c)); /* Dummy for logic which depends on auxiliary data */

171:   /* Assemble System */
172:   PetscCall(DMCreateMatrix(dmSol, &A));
173:   PetscCall(DMCreateGlobalVector(dmSol, &rhs));
174:   PetscCall(DMGetLocalVector(dmForcing, &fLocal));
175:   PetscCall(DMGlobalToLocal(dmForcing, f, INSERT_VALUES, fLocal));

177:   /* Note: if iterating over all the elements, you will usually need to do something
178:      special at one of the boundaries. You can either make use of the existence
179:      of a "extra" partial dummy element on the right/top/front, or you can use a different stencil.
180:      The construction of the reference solution above uses the first method,
181:      so here we will use the second */

183:   PetscCall(DMStagGetIsLastRank(dmSol, &isLastRank, NULL, NULL));
184:   PetscCall(DMStagGetIsFirstRank(dmSol, &isFirstRank, NULL, NULL));
185:   for (e = start; e < start + n; ++e) {
186:     DMStagStencil pos[3];
187:     PetscScalar   val[3];
188:     PetscInt      idxLoc;

190:     idxLoc          = 0;
191:     pos[idxLoc].i   = e;       /* This element in the 1d ordering */
192:     pos[idxLoc].loc = ELEMENT; /* Element-centered dofs (p) */
193:     pos[idxLoc].c   = 0;       /* Component 0 : first (and only) p dof */
194:     val[idxLoc]     = 0.0;     /* p - u'(x) = 0 */
195:     ++idxLoc;

197:     if (isFirstRank && e == start) {
198:       /* Special case on left boundary */
199:       pos[idxLoc].i   = e;    /* This element in the 1d ordering */
200:       pos[idxLoc].loc = LEFT; /* Left vertex */
201:       pos[idxLoc].c   = 0;
202:       val[idxLoc]     = a; /* u(0) = a */
203:       ++idxLoc;
204:     } else {
205:       PetscScalar fVal;
206:       /* Usual case - deal with velocity on left side of cell
207:          Here, we obtain a value of f (even though it's constant here,
208:          this demonstrates the more-realistic case of a pre-computed coefficient) */
209:       pos[idxLoc].i   = e;    /* This element in the 1d ordering */
210:       pos[idxLoc].loc = LEFT; /* vertex-centered dof (u) */
211:       pos[idxLoc].c   = 0;

213:       PetscCall(DMStagVecGetValuesStencil(dmForcing, fLocal, 1, &pos[idxLoc], &fVal));

215:       val[idxLoc] = fVal; /* p'(x) = f, in interior */
216:       ++idxLoc;
217:     }
218:     if (boundary != DM_BOUNDARY_PERIODIC && isLastRank && e == start + n - 1) {
219:       /* Special case on right boundary (in addition to usual case) */
220:       pos[idxLoc].i   = e; /* This element in the 1d ordering */
221:       pos[idxLoc].loc = RIGHT;
222:       pos[idxLoc].c   = 0;
223:       val[idxLoc]     = b; /* u(1) = b */
224:       ++idxLoc;
225:     }
226:     PetscCall(DMStagVecSetValuesStencil(dmSol, rhs, idxLoc, pos, val, INSERT_VALUES));
227:   }
228:   PetscCall(DMRestoreLocalVector(dmForcing, &fLocal));
229:   PetscCall(VecAssemblyBegin(rhs));
230:   PetscCall(VecAssemblyEnd(rhs));

232:   /* Note: normally it would be more efficient to assemble the RHS and the matrix
233:      in the same loop over elements, but we separate them for clarity here */
234:   PetscCall(DMGetCoordinatesLocal(dmSol, &coordSolLocal));
235:   for (e = start; e < start + n; ++e) {
236:     /* Velocity is either a BC or an interior point */
237:     if (isFirstRank && e == start) {
238:       DMStagStencil row;
239:       PetscScalar   val;

241:       row.i   = e;
242:       row.loc = LEFT;
243:       row.c   = 0;
244:       val     = 1.0;
245:       PetscCall(DMStagMatSetValuesStencil(dmSol, A, 1, &row, 1, &row, &val, INSERT_VALUES));
246:     } else {
247:       DMStagStencil row, col[3];
248:       PetscScalar   val[3], xp[2];

250:       row.i   = e;
251:       row.loc = LEFT; /* In general, opt for LEFT/DOWN/BACK  and iterate over elements */
252:       row.c   = 0;

254:       col[0].i   = e;
255:       col[0].loc = ELEMENT;
256:       col[0].c   = 0;

258:       col[1].i   = e - 1;
259:       col[1].loc = ELEMENT;
260:       col[1].c   = 0;

262:       PetscCall(DMStagVecGetValuesStencil(dmCoordSol, coordSolLocal, 2, col, xp));
263:       h = xp[0] - xp[1];
264:       if (boundary == DM_BOUNDARY_PERIODIC && PetscRealPart(h) < 0.0) h += domainSize;

266:       val[0] = 1.0 / h;
267:       val[1] = -1.0 / h;

269:       /* For convenience, we add an explicit 0 on the diagonal. This is a waste,
270:          but it allows for easier use of a direct solver, if desired */
271:       col[2].i   = e;
272:       col[2].loc = LEFT;
273:       col[2].c   = 0;
274:       val[2]     = 0.0;

276:       PetscCall(DMStagMatSetValuesStencil(dmSol, A, 1, &row, 3, col, val, INSERT_VALUES));
277:     }

279:     /* Additional velocity point (BC) on the right */
280:     if (boundary != DM_BOUNDARY_PERIODIC && isLastRank && e == start + n - 1) {
281:       DMStagStencil row;
282:       PetscScalar   val;

284:       row.i   = e;
285:       row.loc = RIGHT;
286:       row.c   = 0;
287:       val     = 1.0;
288:       PetscCall(DMStagMatSetValuesStencil(dmSol, A, 1, &row, 1, &row, &val, INSERT_VALUES));
289:     }

291:     /* Equation on pressure (element) variables */
292:     {
293:       DMStagStencil row, col[3];
294:       PetscScalar   val[3], xu[2];

296:       row.i   = e;
297:       row.loc = ELEMENT;
298:       row.c   = 0;

300:       col[0].i   = e;
301:       col[0].loc = RIGHT;
302:       col[0].c   = 0;

304:       col[1].i   = e;
305:       col[1].loc = LEFT;
306:       col[1].c   = 0;

308:       PetscCall(DMStagVecGetValuesStencil(dmCoordSol, coordSolLocal, 2, col, xu));
309:       h = xu[0] - xu[1];
310:       if (boundary == DM_BOUNDARY_PERIODIC && PetscRealPart(h) < 0.0) h += domainSize;

312:       val[0] = -1.0 / h;
313:       val[1] = 1.0 / h;

315:       col[2].i   = e;
316:       col[2].loc = ELEMENT;
317:       col[2].c   = 0;
318:       val[2]     = 1.0;

320:       PetscCall(DMStagMatSetValuesStencil(dmSol, A, 1, &row, 3, col, val, INSERT_VALUES));
321:     }
322:   }
323:   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
324:   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));

326:   /* Solve */
327:   PetscCall(DMCreateGlobalVector(dmSol, &sol));
328:   PetscCall(KSPCreate(PETSC_COMM_WORLD, &ksp));
329:   PetscCall(KSPSetOperators(ksp, A, A));
330:   PetscCall(KSPGetPC(ksp, &pc));
331:   PetscCall(PCSetType(pc, PCJACOBI)); /* A simple, but non-scalable, solver choice */
332:   PetscCall(KSPSetFromOptions(ksp));
333:   PetscCall(KSPSolve(ksp, rhs, sol));

335:   /* View the components of the solution, demonstrating DMStagMigrateVec() */
336:   {
337:     DM  dmVertsOnly, dmElementsOnly;
338:     Vec u, p;

340:     PetscCall(DMStagCreateCompatibleDMStag(dmSol, 1, 0, 0, 0, &dmVertsOnly));
341:     PetscCall(DMStagCreateCompatibleDMStag(dmSol, 0, 1, 0, 0, &dmElementsOnly));
342:     PetscCall(DMGetGlobalVector(dmVertsOnly, &u));
343:     PetscCall(DMGetGlobalVector(dmElementsOnly, &p));

345:     PetscCall(DMStagMigrateVec(dmSol, sol, dmVertsOnly, u));
346:     PetscCall(DMStagMigrateVec(dmSol, sol, dmElementsOnly, p));

348:     PetscCall(PetscObjectSetName((PetscObject)u, "Sol_u"));
349:     PetscCall(VecView(u, PETSC_VIEWER_STDOUT_WORLD));
350:     PetscCall(PetscObjectSetName((PetscObject)p, "Sol_p"));
351:     PetscCall(VecView(p, PETSC_VIEWER_STDOUT_WORLD));

353:     PetscCall(DMRestoreGlobalVector(dmVertsOnly, &u));
354:     PetscCall(DMRestoreGlobalVector(dmElementsOnly, &p));
355:     PetscCall(DMDestroy(&dmVertsOnly));
356:     PetscCall(DMDestroy(&dmElementsOnly));
357:   }

359:   /* Check Solution */
360:   {
361:     Vec       diff;
362:     PetscReal normsolRef, errAbs, errRel;

364:     PetscCall(VecDuplicate(sol, &diff));
365:     PetscCall(VecCopy(sol, diff));
366:     PetscCall(VecAXPY(diff, -1.0, solRef));
367:     PetscCall(VecNorm(diff, NORM_2, &errAbs));
368:     PetscCall(VecNorm(solRef, NORM_2, &normsolRef));
369:     errRel = errAbs / normsolRef;
370:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Error (abs): %g\nError (rel): %g\n", (double)errAbs, (double)errRel));
371:     PetscCall(VecDestroy(&diff));
372:   }

374:   /* Clean up and finalize PETSc */
375:   PetscCall(KSPDestroy(&ksp));
376:   PetscCall(VecDestroy(&sol));
377:   PetscCall(VecDestroy(&solRef));
378:   PetscCall(VecDestroy(&rhs));
379:   PetscCall(VecDestroy(&f));
380:   PetscCall(MatDestroy(&A));
381:   PetscCall(DMDestroy(&dmSol));
382:   PetscCall(DMDestroy(&dmForcing));
383:   PetscCall(PetscFinalize());
384:   return 0;
385: }

387: /*TEST

389:    test:
390:       suffix: 1
391:       nsize: 7
392:       args: -dm_view -stag_grid_x 11 -stag_stencil_type star -a 1.33 -b 7.22 -c 347.2 -ksp_monitor_short

394:    test:
395:       suffix: periodic
396:       nsize: 3
397:       args: -dm_view -stag_grid_x 13 -stag_boundary_type_x periodic -a 1.1234

399:    test:
400:       suffix: periodic_seq
401:       nsize: 1
402:       args: -dm_view -stag_grid_x 13 -stag_boundary_type_x periodic -a 1.1234

404:    test:
405:       suffix: ghosted_vacuous
406:       nsize: 3
407:       args: -dm_view -stag_grid_x 13 -stag_boundary_type_x ghosted -a 1.1234

409: TEST*/
```