Actual source code: stagda.c
1: /* Routines to convert between a (subset of) DMStag and DMDA */
3: #include <petscdmda.h>
4: #include <petsc/private/dmstagimpl.h>
5: #include <petscdmdatypes.h>
7: static PetscErrorCode DMStagCreateCompatibleDMDA(DM dm, DMStagStencilLocation loc, PetscInt c, DM *dmda)
8: {
9: DM_Stag *const stag = (DM_Stag *)dm->data;
10: PetscInt dim, i, j, stencilWidth, dof, N[DMSTAG_MAX_DIM];
11: DMDAStencilType stencilType;
12: PetscInt *l[DMSTAG_MAX_DIM];
15: DMGetDimension(dm, &dim);
17: /* Create grid decomposition (to be adjusted later) */
18: for (i = 0; i < dim; ++i) {
19: PetscMalloc1(stag->nRanks[i], &l[i]);
20: for (j = 0; j < stag->nRanks[i]; ++j) l[i][j] = stag->l[i][j];
21: N[i] = stag->N[i];
22: }
24: /* dof */
25: dof = c < 0 ? -c : 1;
27: /* Determine/adjust sizes */
28: switch (loc) {
29: case DMSTAG_ELEMENT:
30: break;
31: case DMSTAG_LEFT:
32: case DMSTAG_RIGHT:
34: l[0][stag->nRanks[0] - 1] += 1; /* extra vertex in direction 0 on last rank in dimension 0 */
35: N[0] += 1;
36: break;
37: case DMSTAG_UP:
38: case DMSTAG_DOWN:
40: l[1][stag->nRanks[1] - 1] += 1; /* extra vertex in direction 1 on last rank in dimension 1 */
41: N[1] += 1;
42: break;
43: case DMSTAG_BACK:
44: case DMSTAG_FRONT:
46: l[2][stag->nRanks[2] - 1] += 1; /* extra vertex in direction 2 on last rank in dimension 2 */
47: N[2] += 1;
48: break;
49: case DMSTAG_DOWN_LEFT:
50: case DMSTAG_DOWN_RIGHT:
51: case DMSTAG_UP_LEFT:
52: case DMSTAG_UP_RIGHT:
54: for (i = 0; i < 2; ++i) { /* extra vertex in direction i on last rank in dimension i = 0,1 */
55: l[i][stag->nRanks[i] - 1] += 1;
56: N[i] += 1;
57: }
58: break;
59: case DMSTAG_BACK_LEFT:
60: case DMSTAG_BACK_RIGHT:
61: case DMSTAG_FRONT_LEFT:
62: case DMSTAG_FRONT_RIGHT:
64: for (i = 0; i < 3; i += 2) { /* extra vertex in direction i on last rank in dimension i = 0,2 */
65: l[i][stag->nRanks[i] - 1] += 1;
66: N[i] += 1;
67: }
68: break;
69: case DMSTAG_BACK_DOWN:
70: case DMSTAG_BACK_UP:
71: case DMSTAG_FRONT_DOWN:
72: case DMSTAG_FRONT_UP:
74: for (i = 1; i < 3; ++i) { /* extra vertex in direction i on last rank in dimension i = 1,2 */
75: l[i][stag->nRanks[i] - 1] += 1;
76: N[i] += 1;
77: }
78: break;
79: case DMSTAG_BACK_DOWN_LEFT:
80: case DMSTAG_BACK_DOWN_RIGHT:
81: case DMSTAG_BACK_UP_LEFT:
82: case DMSTAG_BACK_UP_RIGHT:
83: case DMSTAG_FRONT_DOWN_LEFT:
84: case DMSTAG_FRONT_DOWN_RIGHT:
85: case DMSTAG_FRONT_UP_LEFT:
86: case DMSTAG_FRONT_UP_RIGHT:
88: for (i = 0; i < 3; ++i) { /* extra vertex in direction i on last rank in dimension i = 0,1,2 */
89: l[i][stag->nRanks[i] - 1] += 1;
90: N[i] += 1;
91: }
92: break;
93: default:
94: SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Not implemented for location %s", DMStagStencilLocations[loc]);
95: }
97: /* Use the same stencil type */
98: switch (stag->stencilType) {
99: case DMSTAG_STENCIL_STAR:
100: stencilType = DMDA_STENCIL_STAR;
101: stencilWidth = stag->stencilWidth;
102: break;
103: case DMSTAG_STENCIL_BOX:
104: stencilType = DMDA_STENCIL_BOX;
105: stencilWidth = stag->stencilWidth;
106: break;
107: default:
108: SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Unsupported Stencil Type %d", stag->stencilType);
109: }
111: /* Create DMDA, using same boundary type */
112: switch (dim) {
113: case 1:
114: DMDACreate1d(PetscObjectComm((PetscObject)dm), stag->boundaryType[0], N[0], dof, stencilWidth, l[0], dmda);
115: break;
116: case 2:
117: DMDACreate2d(PetscObjectComm((PetscObject)dm), stag->boundaryType[0], stag->boundaryType[1], stencilType, N[0], N[1], stag->nRanks[0], stag->nRanks[1], dof, stencilWidth, l[0], l[1], dmda);
118: break;
119: case 3:
120: DMDACreate3d(PetscObjectComm((PetscObject)dm), stag->boundaryType[0], stag->boundaryType[1], stag->boundaryType[2], stencilType, N[0], N[1], N[2], stag->nRanks[0], stag->nRanks[1], stag->nRanks[2], dof, stencilWidth, l[0], l[1], l[2], dmda);
121: break;
122: default:
123: SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "not implemented for dim %" PetscInt_FMT, dim);
124: }
125: for (i = 0; i < dim; ++i) PetscFree(l[i]);
126: return 0;
127: }
129: /*
130: Helper function to get the number of extra points in a DMDA representation for a given canonical location.
131: */
132: static PetscErrorCode DMStagDMDAGetExtraPoints(DM dm, DMStagStencilLocation locCanonical, PetscInt *extraPoint)
133: {
134: PetscInt dim, d, nExtra[DMSTAG_MAX_DIM];
137: DMGetDimension(dm, &dim);
139: DMStagGetCorners(dm, NULL, NULL, NULL, NULL, NULL, NULL, &nExtra[0], &nExtra[1], &nExtra[2]);
140: for (d = 0; d < dim; ++d) extraPoint[d] = 0;
141: switch (locCanonical) {
142: case DMSTAG_ELEMENT:
143: break; /* no extra points */
144: case DMSTAG_LEFT:
145: extraPoint[0] = nExtra[0];
146: break; /* only extra point in x */
147: case DMSTAG_DOWN:
148: extraPoint[1] = nExtra[1];
149: break; /* only extra point in y */
150: case DMSTAG_BACK:
151: extraPoint[2] = nExtra[2];
152: break; /* only extra point in z */
153: case DMSTAG_DOWN_LEFT:
154: extraPoint[0] = nExtra[0];
155: extraPoint[1] = nExtra[1];
156: break; /* extra point in both x and y */
157: case DMSTAG_BACK_LEFT:
158: extraPoint[0] = nExtra[0];
159: extraPoint[2] = nExtra[2];
160: break; /* extra point in both x and z */
161: case DMSTAG_BACK_DOWN:
162: extraPoint[1] = nExtra[1];
163: extraPoint[2] = nExtra[2];
164: break; /* extra point in both y and z */
165: case DMSTAG_BACK_DOWN_LEFT:
166: extraPoint[0] = nExtra[0];
167: extraPoint[1] = nExtra[1];
168: extraPoint[2] = nExtra[2];
169: break; /* extra points in x,y,z */
170: default:
171: SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Not implemented for location (perhaps not canonical) %s", DMStagStencilLocations[locCanonical]);
172: }
173: return 0;
174: }
176: /*
177: Function much like DMStagMigrateVec(), but which accepts an additional position argument to disambiguate which
178: type of DMDA to migrate to.
179: */
181: static PetscErrorCode DMStagMigrateVecDMDA(DM dm, Vec vec, DMStagStencilLocation loc, PetscInt c, DM dmTo, Vec vecTo)
182: {
183: PetscInt i, j, k, d, dim, dof, dofToMax, start[DMSTAG_MAX_DIM], n[DMSTAG_MAX_DIM], extraPoint[DMSTAG_MAX_DIM];
184: Vec vecLocal;
190: DMGetDimension(dm, &dim);
191: DMDAGetDof(dmTo, &dofToMax);
193: DMStagGetCorners(dm, &start[0], &start[1], &start[2], &n[0], &n[1], &n[2], NULL, NULL, NULL);
194: DMStagDMDAGetExtraPoints(dm, loc, extraPoint);
195: DMStagGetLocationDOF(dm, loc, &dof);
196: DMGetLocalVector(dm, &vecLocal);
197: DMGlobalToLocalBegin(dm, vec, INSERT_VALUES, vecLocal);
198: DMGlobalToLocalEnd(dm, vec, INSERT_VALUES, vecLocal);
199: if (dim == 1) {
200: PetscScalar **arrTo;
201: DMDAVecGetArrayDOF(dmTo, vecTo, &arrTo);
202: if (c < 0) {
203: const PetscInt dofTo = -c;
204: for (i = start[0]; i < start[0] + n[0] + extraPoint[0]; ++i) {
205: for (d = 0; d < PetscMin(dof, dofTo); ++d) {
206: DMStagStencil pos;
207: pos.i = i;
208: pos.loc = loc;
209: pos.c = d;
210: DMStagVecGetValuesStencil(dm, vecLocal, 1, &pos, &arrTo[i][d]);
211: }
212: for (; d < dofTo; ++d) { arrTo[i][d] = 0.0; /* Pad extra dof with zeros */ }
213: }
214: } else {
215: for (i = start[0]; i < start[0] + n[0] + extraPoint[0]; ++i) {
216: DMStagStencil pos;
217: pos.i = i;
218: pos.loc = loc;
219: pos.c = c;
220: DMStagVecGetValuesStencil(dm, vecLocal, 1, &pos, &arrTo[i][0]);
221: }
222: }
223: DMDAVecRestoreArrayDOF(dmTo, vecTo, &arrTo);
224: } else if (dim == 2) {
225: PetscScalar ***arrTo;
226: DMDAVecGetArrayDOF(dmTo, vecTo, &arrTo);
227: if (c < 0) {
228: const PetscInt dofTo = -c;
229: for (j = start[1]; j < start[1] + n[1] + extraPoint[1]; ++j) {
230: for (i = start[0]; i < start[0] + n[0] + extraPoint[0]; ++i) {
231: for (d = 0; d < PetscMin(dof, dofTo); ++d) {
232: DMStagStencil pos;
233: pos.i = i;
234: pos.j = j;
235: pos.loc = loc;
236: pos.c = d;
237: DMStagVecGetValuesStencil(dm, vecLocal, 1, &pos, &arrTo[j][i][d]);
238: }
239: for (; d < dofTo; ++d) { arrTo[j][i][d] = 0.0; /* Pad extra dof with zeros */ }
240: }
241: }
242: } else {
243: for (j = start[1]; j < start[1] + n[1] + extraPoint[1]; ++j) {
244: for (i = start[0]; i < start[0] + n[0] + extraPoint[0]; ++i) {
245: DMStagStencil pos;
246: pos.i = i;
247: pos.j = j;
248: pos.loc = loc;
249: pos.c = c;
250: DMStagVecGetValuesStencil(dm, vecLocal, 1, &pos, &arrTo[j][i][0]);
251: }
252: }
253: }
254: DMDAVecRestoreArrayDOF(dmTo, vecTo, &arrTo);
255: } else if (dim == 3) {
256: PetscScalar ****arrTo;
257: DMDAVecGetArrayDOF(dmTo, vecTo, &arrTo);
258: if (c < 0) {
259: const PetscInt dofTo = -c;
260: for (k = start[2]; k < start[2] + n[2] + extraPoint[2]; ++k) {
261: for (j = start[1]; j < start[1] + n[1] + extraPoint[1]; ++j) {
262: for (i = start[0]; i < start[0] + n[0] + extraPoint[0]; ++i) {
263: for (d = 0; d < PetscMin(dof, dofTo); ++d) {
264: DMStagStencil pos;
265: pos.i = i;
266: pos.j = j;
267: pos.k = k;
268: pos.loc = loc;
269: pos.c = d;
270: DMStagVecGetValuesStencil(dm, vecLocal, 1, &pos, &arrTo[k][j][i][d]);
271: }
272: for (; d < dofTo; ++d) { arrTo[k][j][i][d] = 0.0; /* Pad extra dof with zeros */ }
273: }
274: }
275: }
276: } else {
277: for (k = start[2]; k < start[2] + n[2] + extraPoint[2]; ++k) {
278: for (j = start[1]; j < start[1] + n[1] + extraPoint[1]; ++j) {
279: for (i = start[0]; i < start[0] + n[0] + extraPoint[0]; ++i) {
280: DMStagStencil pos;
281: pos.i = i;
282: pos.j = j;
283: pos.k = k;
284: pos.loc = loc;
285: pos.c = c;
286: DMStagVecGetValuesStencil(dm, vecLocal, 1, &pos, &arrTo[k][j][i][0]);
287: }
288: }
289: }
290: }
291: DMDAVecRestoreArrayDOF(dmTo, vecTo, &arrTo);
292: } else SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Unsupported dimension %" PetscInt_FMT, dim);
293: DMRestoreLocalVector(dm, &vecLocal);
294: return 0;
295: }
297: /* Transfer coordinates from a DMStag to a DMDA, specifying which location */
298: static PetscErrorCode DMStagTransferCoordinatesToDMDA(DM dmstag, DMStagStencilLocation loc, DM dmda)
299: {
300: PetscInt dim, start[DMSTAG_MAX_DIM], n[DMSTAG_MAX_DIM], extraPoint[DMSTAG_MAX_DIM], d;
301: DM dmstagCoord, dmdaCoord;
302: DMType dmstagCoordType;
303: Vec stagCoord, daCoord;
304: PetscBool daCoordIsStag, daCoordIsProduct;
308: DMGetDimension(dmstag, &dim);
309: DMGetCoordinateDM(dmstag, &dmstagCoord);
310: DMGetCoordinatesLocal(dmstag, &stagCoord); /* Note local */
311: DMGetCoordinateDM(dmda, &dmdaCoord);
312: daCoord = NULL;
313: DMGetCoordinates(dmda, &daCoord);
314: if (!daCoord) {
315: DMCreateGlobalVector(dmdaCoord, &daCoord);
316: DMSetCoordinates(dmda, daCoord);
317: VecDestroy(&daCoord);
318: DMGetCoordinates(dmda, &daCoord);
319: }
320: DMGetType(dmstagCoord, &dmstagCoordType);
321: PetscStrcmp(dmstagCoordType, DMSTAG, &daCoordIsStag);
322: PetscStrcmp(dmstagCoordType, DMPRODUCT, &daCoordIsProduct);
323: DMStagGetCorners(dmstag, &start[0], &start[1], &start[2], &n[0], &n[1], &n[2], NULL, NULL, NULL);
324: DMStagDMDAGetExtraPoints(dmstag, loc, extraPoint);
325: if (dim == 1) {
326: PetscInt ex;
327: PetscScalar **cArrDa;
328: DMDAVecGetArrayDOF(dmdaCoord, daCoord, &cArrDa);
329: if (daCoordIsStag) {
330: PetscInt slot;
331: PetscScalar **cArrStag;
332: DMStagGetLocationSlot(dmstagCoord, loc, 0, &slot);
333: DMStagVecGetArrayRead(dmstagCoord, stagCoord, &cArrStag);
334: for (ex = start[0]; ex < start[0] + n[0] + extraPoint[0]; ++ex) cArrDa[ex][0] = cArrStag[ex][slot];
335: DMStagVecRestoreArrayRead(dmstagCoord, stagCoord, &cArrStag);
336: } else if (daCoordIsProduct) {
337: PetscScalar **cArrX;
338: DMStagGetProductCoordinateArraysRead(dmstag, &cArrX, NULL, NULL);
339: for (ex = start[0]; ex < start[0] + n[0] + extraPoint[0]; ++ex) cArrDa[ex][0] = cArrX[ex][0];
340: DMStagRestoreProductCoordinateArraysRead(dmstag, &cArrX, NULL, NULL);
341: } else SETERRQ(PetscObjectComm((PetscObject)dmstag), PETSC_ERR_SUP, "Stag to DA coordinate transfer only supported for DMStag coordinate DM of type DMstag or DMProduct");
342: DMDAVecRestoreArrayDOF(dmdaCoord, daCoord, &cArrDa);
343: } else if (dim == 2) {
344: PetscInt ex, ey;
345: PetscScalar ***cArrDa;
346: DMDAVecGetArrayDOF(dmdaCoord, daCoord, &cArrDa);
347: if (daCoordIsStag) {
348: PetscInt slot;
349: PetscScalar ***cArrStag;
350: DMStagGetLocationSlot(dmstagCoord, loc, 0, &slot);
351: DMStagVecGetArrayRead(dmstagCoord, stagCoord, &cArrStag);
352: for (ey = start[1]; ey < start[1] + n[1] + extraPoint[1]; ++ey) {
353: for (ex = start[0]; ex < start[0] + n[0] + extraPoint[0]; ++ex) {
354: for (d = 0; d < 2; ++d) cArrDa[ey][ex][d] = cArrStag[ey][ex][slot + d];
355: }
356: }
357: DMStagVecRestoreArrayRead(dmstagCoord, stagCoord, &cArrStag);
358: } else if (daCoordIsProduct) {
359: PetscScalar **cArrX, **cArrY;
360: DMStagGetProductCoordinateArraysRead(dmstag, &cArrX, &cArrY, NULL);
361: for (ey = start[1]; ey < start[1] + n[1] + extraPoint[1]; ++ey) {
362: for (ex = start[0]; ex < start[0] + n[0] + extraPoint[0]; ++ex) {
363: cArrDa[ey][ex][0] = cArrX[ex][0];
364: cArrDa[ey][ex][1] = cArrY[ey][0];
365: }
366: }
367: DMStagRestoreProductCoordinateArraysRead(dmstag, &cArrX, &cArrY, NULL);
368: } else SETERRQ(PetscObjectComm((PetscObject)dmstag), PETSC_ERR_SUP, "Stag to DA coordinate transfer only supported for DMStag coordinate DM of type DMstag or DMProduct");
369: DMDAVecRestoreArrayDOF(dmdaCoord, daCoord, &cArrDa);
370: } else if (dim == 3) {
371: PetscInt ex, ey, ez;
372: PetscScalar ****cArrDa;
373: DMDAVecGetArrayDOF(dmdaCoord, daCoord, &cArrDa);
374: if (daCoordIsStag) {
375: PetscInt slot;
376: PetscScalar ****cArrStag;
377: DMStagGetLocationSlot(dmstagCoord, loc, 0, &slot);
378: DMStagVecGetArrayRead(dmstagCoord, stagCoord, &cArrStag);
379: for (ez = start[2]; ez < start[2] + n[2] + extraPoint[2]; ++ez) {
380: for (ey = start[1]; ey < start[1] + n[1] + extraPoint[1]; ++ey) {
381: for (ex = start[0]; ex < start[0] + n[0] + extraPoint[0]; ++ex) {
382: for (d = 0; d < 3; ++d) cArrDa[ez][ey][ex][d] = cArrStag[ez][ey][ex][slot + d];
383: }
384: }
385: }
386: DMStagVecRestoreArrayRead(dmstagCoord, stagCoord, &cArrStag);
387: } else if (daCoordIsProduct) {
388: PetscScalar **cArrX, **cArrY, **cArrZ;
389: DMStagGetProductCoordinateArraysRead(dmstag, &cArrX, &cArrY, &cArrZ);
390: for (ez = start[2]; ez < start[2] + n[2] + extraPoint[2]; ++ez) {
391: for (ey = start[1]; ey < start[1] + n[1] + extraPoint[1]; ++ey) {
392: for (ex = start[0]; ex < start[0] + n[0] + extraPoint[0]; ++ex) {
393: cArrDa[ez][ey][ex][0] = cArrX[ex][0];
394: cArrDa[ez][ey][ex][1] = cArrY[ey][0];
395: cArrDa[ez][ey][ex][2] = cArrZ[ez][0];
396: }
397: }
398: }
399: DMStagRestoreProductCoordinateArraysRead(dmstag, &cArrX, &cArrY, &cArrZ);
400: } else SETERRQ(PetscObjectComm((PetscObject)dmstag), PETSC_ERR_SUP, "Stag to DA coordinate transfer only supported for DMStag coordinate DM of type DMstag or DMProduct");
401: DMDAVecRestoreArrayDOF(dmdaCoord, daCoord, &cArrDa);
402: } else SETERRQ(PetscObjectComm((PetscObject)dmstag), PETSC_ERR_SUP, "Unsupported dimension %" PetscInt_FMT, dim);
403: return 0;
404: }
406: /*@C
407: DMStagVecSplitToDMDA - create a `DMDA` and `Vec` from a subgrid of a `DMSTAG` and its `Vec`
409: Collective
411: Input Parameters:
412: + dm - the `DMSTAG` object
413: . vec- Vec object associated with `dm`
414: . loc - which subgrid to extract (see `DMStagStencilLocation`)
415: - c - which component to extract (see note below)
417: Output Parameters:
418: + pda - the `DMDA`
419: - pdavec - the new `Vec`
421: Level: advanced
423: Notes:
424: If a `c` value of `-k` is provided, the first `k` DOF for that position are extracted,
425: padding with zero values if needed. If a non-negative value is provided, a single
426: DOF is extracted.
428: The caller is responsible for destroying the created `DMDA` and `Vec`.
430: .seealso: [](chapter_stag), `DMSTAG`, `DMDA`, `DMStagStencilLocation`, `DM`, `Vec`, `DMStagMigrateVec()`, `DMStagCreateCompatibleDMStag()`
431: @*/
432: PetscErrorCode DMStagVecSplitToDMDA(DM dm, Vec vec, DMStagStencilLocation loc, PetscInt c, DM *pda, Vec *pdavec)
433: {
434: PetscInt dim, locdof;
435: DM da, coordDM;
436: Vec davec;
437: DMStagStencilLocation locCanonical;
441: DMGetDimension(dm, &dim);
442: DMStagGetLocationDOF(dm, loc, &locdof);
444: DMStagStencilLocationCanonicalize(loc, &locCanonical);
445: DMStagCreateCompatibleDMDA(dm, locCanonical, c, pda);
446: da = *pda;
447: DMSetUp(*pda);
448: if (dm->coordinates[0].dm != NULL) {
449: DMGetCoordinateDM(dm, &coordDM);
450: DMStagTransferCoordinatesToDMDA(dm, locCanonical, da);
451: }
452: DMCreateGlobalVector(da, pdavec);
453: davec = *pdavec;
454: DMStagMigrateVecDMDA(dm, vec, locCanonical, c, da, davec);
455: return 0;
456: }