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:   PetscErrorCode    ierr;
85:   DM                dmSol,dmForcing;
86:   DM                dmCoordSol;
87:   Vec               sol,solRef,solRefLocal,f,fLocal,rhs,coordSolLocal;
88:   Mat               A;
89:   PetscScalar       a,b,c,h;
90:   KSP               ksp;
91:   PC                pc;
92:   PetscInt          start,n,e,nExtra;
93:   PetscInt          iu,ip,ixu,ixp;
94:   PetscBool         isLastRank,isFirstRank;
95:   PetscScalar       **arrSol,**arrCoordSol;
96:   DMBoundaryType    boundary;

98:   const PetscReal domainSize = 1.0;

100:   /* Initialize PETSc */
101:   PetscInitialize(&argc,&argv,(char*)0,help);if (ierr) return ierr;

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:   DMStagCreate1d(PETSC_COMM_WORLD,DM_BOUNDARY_NONE,3,1,1,DMSTAG_STENCIL_BOX,1,NULL,&dmSol);
107:   DMSetFromOptions(dmSol);
108:   DMSetUp(dmSol);

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

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

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

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

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

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

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

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

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

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

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

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

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

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

230:   /* Note: normally it would be more efficient to assemble the RHS and the matrix
231:      in the same loop over elements, but we separate them for clarity here */
232:   DMGetCoordinatesLocal(dmSol,&coordSolLocal);
233:   for (e = start; e<start+n; ++e) {

235:     /* Velocity is either a BC or an interior point */
236:     if (isFirstRank && e == start) {
237:       DMStagStencil row;
238:       PetscScalar   val;

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

386: /*TEST

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

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

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

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

408: TEST*/
```