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 /*
6
7 To demonstrate the basic functionality of DMStag, solves a second-order ODE,
8
9 u''(x) = f(x), 0 < x < 1
10 u(0) = a
11 u(1) = b
12
13 in mixed form, that is by introducing an auxiliary variable p
14
15 p'(x) = f(x), p - u'(x) = 0, 0 < x < 1
16 u(0) = a
17 u(1) = b
18
19 For f == c, the solution is
20
21 u(x) = a + (b - a - (c/2)) x + (c/2) x^2
22 p(x) = (b - a - (c/2)) + c x
23
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.
27
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.
32
33 Parameters for the main grid can be controlled with command line options, e.g.
34
35 -stag_grid_x 10
36
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().
40
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.
46
47 DMStagStencil stencil;
48 stencil.i = i
49 stencil.loc = DMSTAG_LEFT;
50 stencil.c = 0
51
52 Note that below, for convenenience, we #define an alias LEFT for DMSTAG_LEFT.
53
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().
58
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
63 direct solver, e.g. configure --download-suitesparse and run
64
65 $PETSC_DIR/$PETSC_ARCH/bin/mpiexec -n 3 ./stag_ex2 -stag_grid_x 100 -pc_type lu -pc_factor_mat_solver_package umfpack
66
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.
69
70 -stag_boundary_type_x periodic
71
72 */
73 #include <petscdm.h>
74 #include <petscksp.h>
75 #include <petscdmstag.h>
76
77 /* Shorter, more convenient names for DMStagStencilLocation entries */
78 #define LEFT DMSTAG_LEFT
79 #define RIGHT DMSTAG_RIGHT
80 #define ELEMENT DMSTAG_ELEMENT
81
main(int argc,char ** argv)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;
96
97 const PetscReal domainSize = 1.0;
98
99 /* Initialize PETSc */
100 PetscFunctionBeginUser;
101 PetscCall(PetscInitialize(&argc, &argv, NULL, help));
102
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));
109
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));
113
114 /* Determine boundary type */
115 PetscCall(DMStagGetBoundaryTypes(dmSol, &boundary, NULL, NULL));
116
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 }
129
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));
136 PetscCall(DMStagVecGetArrayRead(dmCoordSol, coordSolLocal, &arrCoordSol));
137 PetscCall(DMStagGetCorners(dmSol, &start, NULL, NULL, &n, NULL, NULL, &nExtra, NULL, NULL));
138
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 }
154 PetscCall(DMStagVecRestoreArrayRead(dmCoordSol, coordSolLocal, &arrCoordSol));
155 PetscCall(DMStagVecRestoreArray(dmSol, solRefLocal, &arrSol));
156 PetscCall(DMLocalToGlobal(dmSol, solRefLocal, INSERT_VALUES, solRef));
157 PetscCall(DMRestoreLocalVector(dmSol, &solRefLocal));
158
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.
162
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 */
170
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));
176
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 */
182
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;
189
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;
196
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;
212
213 PetscCall(DMStagVecGetValuesStencil(dmForcing, fLocal, 1, &pos[idxLoc], &fVal));
214
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));
231
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;
240
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];
249
250 row.i = e;
251 row.loc = LEFT; /* In general, opt for LEFT/DOWN/BACK and iterate over elements */
252 row.c = 0;
253
254 col[0].i = e;
255 col[0].loc = ELEMENT;
256 col[0].c = 0;
257
258 col[1].i = e - 1;
259 col[1].loc = ELEMENT;
260 col[1].c = 0;
261
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;
265
266 val[0] = 1.0 / h;
267 val[1] = -1.0 / h;
268
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;
275
276 PetscCall(DMStagMatSetValuesStencil(dmSol, A, 1, &row, 3, col, val, INSERT_VALUES));
277 }
278
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;
283
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 }
290
291 /* Equation on pressure (element) variables */
292 {
293 DMStagStencil row, col[3];
294 PetscScalar val[3], xu[2];
295
296 row.i = e;
297 row.loc = ELEMENT;
298 row.c = 0;
299
300 col[0].i = e;
301 col[0].loc = RIGHT;
302 col[0].c = 0;
303
304 col[1].i = e;
305 col[1].loc = LEFT;
306 col[1].c = 0;
307
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;
311
312 val[0] = -1.0 / h;
313 val[1] = 1.0 / h;
314
315 col[2].i = e;
316 col[2].loc = ELEMENT;
317 col[2].c = 0;
318 val[2] = 1.0;
319
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));
325
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));
334
335 /* View the components of the solution, demonstrating DMStagMigrateVec() */
336 {
337 DM dmVertsOnly, dmElementsOnly;
338 Vec u, p;
339
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));
344
345 PetscCall(DMStagMigrateVec(dmSol, sol, dmVertsOnly, u));
346 PetscCall(DMStagMigrateVec(dmSol, sol, dmElementsOnly, p));
347
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));
352
353 PetscCall(DMRestoreGlobalVector(dmVertsOnly, &u));
354 PetscCall(DMRestoreGlobalVector(dmElementsOnly, &p));
355 PetscCall(DMDestroy(&dmVertsOnly));
356 PetscCall(DMDestroy(&dmElementsOnly));
357 }
358
359 /* Check Solution */
360 {
361 Vec diff;
362 PetscReal normsolRef, errAbs, errRel;
363
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 }
373
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 }
386
387 /*TEST
388
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
393
394 test:
395 suffix: periodic
396 nsize: 3
397 args: -dm_view -stag_grid_x 13 -stag_boundary_type_x periodic -a 1.1234
398
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
403
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
408
409 TEST*/
410