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: }