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 : SETERRQ(PetscObjectComm((PetscObject)dm),PETSC_ERR_SUP,"Not implemented for location %s",DMStagStencilLocations[loc]);
 94:   }

 96:   /* Use the same stencil type */
 97:   switch (stag->stencilType) {
 98:     case DMSTAG_STENCIL_STAR: stencilType = DMDA_STENCIL_STAR; stencilWidth = stag->stencilWidth; break;
 99:     case DMSTAG_STENCIL_BOX : stencilType = DMDA_STENCIL_BOX ; stencilWidth = stag->stencilWidth; break;
100:     default: SETERRQ(PetscObjectComm((PetscObject)dm),PETSC_ERR_SUP,"Unsupported Stencil Type %d",stag->stencilType);
101:   }

103:   /* Create DMDA, using same boundary type */
104:   switch (dim) {
105:     case 1:
106:       DMDACreate1d(PetscObjectComm((PetscObject)dm),stag->boundaryType[0],N[0],dof,stencilWidth,l[0],dmda);
107:       break;
108:     case 2:
109:       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);
110:       break;
111:     case 3:
112:       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);
113:       break;
114:     default: SETERRQ(PetscObjectComm((PetscObject)dm),PETSC_ERR_SUP,"not implemented for dim %d",dim);
115:   }
116:   for (i=0; i<dim; ++i) {
117:     PetscFree(l[i]);
118:   }
119:   return 0;
120: }

122: /*
123: Helper function to get the number of extra points in a DMDA representation for a given canonical location.
124: */
125: static PetscErrorCode DMStagDMDAGetExtraPoints(DM dm,DMStagStencilLocation locCanonical,PetscInt *extraPoint)
126: {
127:   PetscInt       dim,d,nExtra[DMSTAG_MAX_DIM];

130:   DMGetDimension(dm,&dim);
132:   DMStagGetCorners(dm,NULL,NULL,NULL,NULL,NULL,NULL,&nExtra[0],&nExtra[1],&nExtra[2]);
133:   for (d=0; d<dim; ++d) extraPoint[d] = 0;
134:   switch (locCanonical) {
135:     case DMSTAG_ELEMENT:
136:       break; /* no extra points */
137:     case DMSTAG_LEFT:
138:       extraPoint[0] = nExtra[0]; break; /* only extra point in x */
139:     case DMSTAG_DOWN:
140:       extraPoint[1] = nExtra[1]; break; /* only extra point in y */
141:     case DMSTAG_BACK:
142:       extraPoint[2] = nExtra[2]; break; /* only extra point in z */
143:     case DMSTAG_DOWN_LEFT:
144:       extraPoint[0] = nExtra[0]; extraPoint[1] = nExtra[1]; break; /* extra point in both x and y  */
145:     case DMSTAG_BACK_LEFT:
146:       extraPoint[0] = nExtra[0]; extraPoint[2] = nExtra[2]; break; /* extra point in both x and z  */
147:     case DMSTAG_BACK_DOWN:
148:       extraPoint[1] = nExtra[1]; extraPoint[2] = nExtra[2]; break; /* extra point in both y and z  */
149:     case DMSTAG_BACK_DOWN_LEFT:
150:      extraPoint[0] = nExtra[0]; extraPoint[1] = nExtra[1]; extraPoint[2] = nExtra[2]; break; /* extra points in x,y,z */
151:     default : SETERRQ(PetscObjectComm((PetscObject)dm),PETSC_ERR_SUP,"Not implemented for location (perhaps not canonical) %s",DMStagStencilLocations[locCanonical]);
152:   }
153:   return 0;
154: }

156: /*
157: Function much like DMStagMigrateVec(), but which accepts an additional position argument to disambiguate which
158: type of DMDA to migrate to.
159: */

161: static PetscErrorCode DMStagMigrateVecDMDA(DM dm,Vec vec,DMStagStencilLocation loc,PetscInt c,DM dmTo,Vec vecTo)
162: {
163:   PetscInt       i,j,k,d,dim,dof,dofToMax,start[DMSTAG_MAX_DIM],n[DMSTAG_MAX_DIM],extraPoint[DMSTAG_MAX_DIM];
164:   Vec            vecLocal;

170:   DMGetDimension(dm,&dim);
171:   DMDAGetDof(dmTo,&dofToMax);
173:   DMStagGetCorners(dm,&start[0],&start[1],&start[2],&n[0],&n[1],&n[2],NULL,NULL,NULL);
174:   DMStagDMDAGetExtraPoints(dm,loc,extraPoint);
175:   DMStagGetLocationDOF(dm,loc,&dof);
176:   DMGetLocalVector(dm,&vecLocal);
177:   DMGlobalToLocalBegin(dm,vec,INSERT_VALUES,vecLocal);
178:   DMGlobalToLocalEnd(dm,vec,INSERT_VALUES,vecLocal);
179:   if (dim == 1) {
180:     PetscScalar **arrTo;
181:     DMDAVecGetArrayDOF(dmTo,vecTo,&arrTo);
182:     if (c < 0) {
183:       const PetscInt dofTo = -c;
184:       for (i=start[0]; i<start[0] + n[0] + extraPoint[0]; ++i) {
185:         for (d=0; d<PetscMin(dof,dofTo); ++d) {
186:           DMStagStencil pos;
187:           pos.i = i; pos.loc = loc; pos.c = d;
188:           DMStagVecGetValuesStencil(dm,vecLocal,1,&pos,&arrTo[i][d]);
189:         }
190:         for (;d<dofTo; ++d) {
191:           arrTo[i][d] = 0.0; /* Pad extra dof with zeros */
192:         }
193:       }
194:     } else {
195:       for (i=start[0]; i<start[0] + n[0] + extraPoint[0]; ++i) {
196:         DMStagStencil pos;
197:         pos.i = i; pos.loc = loc; pos.c = c;
198:         DMStagVecGetValuesStencil(dm,vecLocal,1,&pos,&arrTo[i][0]);
199:       }
200:     }
201:     DMDAVecRestoreArrayDOF(dmTo,vecTo,&arrTo);
202:   } else if (dim == 2) {
203:     PetscScalar ***arrTo;
204:     DMDAVecGetArrayDOF(dmTo,vecTo,&arrTo);
205:     if (c < 0) {
206:       const PetscInt dofTo = -c;
207:       for (j=start[1]; j<start[1] + n[1] + extraPoint[1]; ++j) {
208:         for (i=start[0]; i<start[0] + n[0] + extraPoint[0]; ++i) {
209:           for (d=0; d<PetscMin(dof,dofTo); ++d) {
210:             DMStagStencil pos;
211:             pos.i = i; pos.j = j; pos.loc = loc; pos.c = d;
212:             DMStagVecGetValuesStencil(dm,vecLocal,1,&pos,&arrTo[j][i][d]);
213:           }
214:           for (;d<dofTo; ++d) {
215:             arrTo[j][i][d] = 0.0; /* Pad extra dof with zeros */
216:           }
217:         }
218:       }
219:     } else {
220:       for (j=start[1]; j<start[1] + n[1] + extraPoint[1]; ++j) {
221:         for (i=start[0]; i<start[0] + n[0] + extraPoint[0]; ++i) {
222:           DMStagStencil pos;
223:           pos.i = i; pos.j = j; pos.loc = loc; pos.c = c;
224:           DMStagVecGetValuesStencil(dm,vecLocal,1,&pos,&arrTo[j][i][0]);
225:         }
226:       }
227:     }
228:     DMDAVecRestoreArrayDOF(dmTo,vecTo,&arrTo);
229:   } else if (dim == 3) {
230:     PetscScalar ****arrTo;
231:     DMDAVecGetArrayDOF(dmTo,vecTo,&arrTo);
232:     if (c < 0) {
233:       const PetscInt dofTo = -c;
234:       for (k=start[2]; k<start[2] + n[2] + extraPoint[2]; ++k) {
235:         for (j=start[1]; j<start[1] + n[1] + extraPoint[1]; ++j) {
236:           for (i=start[0]; i<start[0] + n[0] + extraPoint[0]; ++i) {
237:             for (d=0; d<PetscMin(dof,dofTo); ++d) {
238:               DMStagStencil pos;
239:               pos.i = i; pos.j = j; pos.k = k; pos.loc = loc; pos.c = d;
240:               DMStagVecGetValuesStencil(dm,vecLocal,1,&pos,&arrTo[k][j][i][d]);
241:             }
242:             for (;d<dofTo; ++d) {
243:               arrTo[k][j][i][d] = 0.0; /* Pad extra dof with zeros */
244:             }
245:           }
246:         }
247:       }
248:     } else {
249:       for (k=start[2]; k<start[2] + n[2] + extraPoint[2]; ++k) {
250:         for (j=start[1]; j<start[1] + n[1] + extraPoint[1]; ++j) {
251:           for (i=start[0]; i<start[0] + n[0] + extraPoint[0]; ++i) {
252:             DMStagStencil pos;
253:             pos.i = i; pos.j = j; pos.k = k; pos.loc = loc; pos.c = c;
254:             DMStagVecGetValuesStencil(dm,vecLocal,1,&pos,&arrTo[k][j][i][0]);
255:           }
256:         }
257:       }
258:     }
259:     DMDAVecRestoreArrayDOF(dmTo,vecTo,&arrTo);
260:   } else SETERRQ(PetscObjectComm((PetscObject)dm),PETSC_ERR_SUP,"Unsupported dimension %d",dim);
261:   DMRestoreLocalVector(dm,&vecLocal);
262:   return 0;
263: }

265: /* Transfer coordinates from a DMStag to a DMDA, specifying which location */
266: static PetscErrorCode DMStagTransferCoordinatesToDMDA(DM dmstag,DMStagStencilLocation loc,DM dmda)
267: {
268:   PetscInt       dim,start[DMSTAG_MAX_DIM],n[DMSTAG_MAX_DIM],extraPoint[DMSTAG_MAX_DIM],d;
269:   DM             dmstagCoord,dmdaCoord;
270:   DMType         dmstagCoordType;
271:   Vec            stagCoord,daCoord;
272:   PetscBool      daCoordIsStag,daCoordIsProduct;

276:   DMGetDimension(dmstag,&dim);
277:   DMGetCoordinateDM(dmstag,&dmstagCoord);
278:   DMGetCoordinatesLocal(dmstag,&stagCoord); /* Note local */
279:   DMGetCoordinateDM(dmda,&dmdaCoord);
280:   daCoord = NULL;
281:   DMGetCoordinates(dmda,&daCoord);
282:   if (!daCoord) {
283:     DMCreateGlobalVector(dmdaCoord,&daCoord);
284:     DMSetCoordinates(dmda,daCoord);
285:     VecDestroy(&daCoord);
286:     DMGetCoordinates(dmda,&daCoord);
287:   }
288:   DMGetType(dmstagCoord,&dmstagCoordType);
289:   PetscStrcmp(dmstagCoordType,DMSTAG,&daCoordIsStag);
290:   PetscStrcmp(dmstagCoordType,DMPRODUCT,&daCoordIsProduct);
291:   DMStagGetCorners(dmstag,&start[0],&start[1],&start[2],&n[0],&n[1],&n[2],NULL,NULL,NULL);
292:   DMStagDMDAGetExtraPoints(dmstag,loc,extraPoint);
293:   if (dim == 1) {
294:     PetscInt ex;
295:     PetscScalar **cArrDa;
296:     DMDAVecGetArrayDOF(dmdaCoord,daCoord,&cArrDa);
297:     if (daCoordIsStag)  {
298:       PetscInt slot;
299:       PetscScalar **cArrStag;
300:       DMStagGetLocationSlot(dmstagCoord,loc,0,&slot);
301:       DMStagVecGetArrayRead(dmstagCoord,stagCoord,&cArrStag);
302:       for (ex=start[0]; ex<start[0] + n[0] + extraPoint[0]; ++ex) {
303:         cArrDa[ex][0] = cArrStag[ex][slot];
304:       }
305:       DMStagVecRestoreArrayRead(dmstagCoord,stagCoord,&cArrStag);
306:     } else if (daCoordIsProduct) {
307:       PetscScalar **cArrX;
308:       DMStagGetProductCoordinateArraysRead(dmstag,&cArrX,NULL,NULL);
309:       for (ex=start[0]; ex<start[0] + n[0] + extraPoint[0]; ++ex) {
310:         cArrDa[ex][0] = cArrX[ex][0];
311:       }
312:       DMStagRestoreProductCoordinateArraysRead(dmstag,&cArrX,NULL,NULL);
313:     } else SETERRQ(PetscObjectComm((PetscObject)dmstag),PETSC_ERR_SUP,"Stag to DA coordinate transfer only supported for DMStag coordinate DM of type DMstag or DMProduct");
314:     DMDAVecRestoreArrayDOF(dmdaCoord,daCoord,&cArrDa);
315:   } else if (dim == 2) {
316:     PetscInt ex,ey;
317:     PetscScalar ***cArrDa;
318:     DMDAVecGetArrayDOF(dmdaCoord,daCoord,&cArrDa);
319:     if (daCoordIsStag)  {
320:       PetscInt slot;
321:       PetscScalar ***cArrStag;
322:       DMStagGetLocationSlot(dmstagCoord,loc,0,&slot);
323:       DMStagVecGetArrayRead(dmstagCoord,stagCoord,&cArrStag);
324:       for (ey=start[1]; ey<start[1] + n[1] + extraPoint[1]; ++ey) {
325:         for (ex=start[0]; ex<start[0] + n[0] + extraPoint[0]; ++ex) {
326:           for (d=0; d<2; ++d) {
327:             cArrDa[ey][ex][d] = cArrStag[ey][ex][slot+d];
328:           }
329:         }
330:       }
331:       DMStagVecRestoreArrayRead(dmstagCoord,stagCoord,&cArrStag);
332:     } else if (daCoordIsProduct) {
333:       PetscScalar **cArrX,**cArrY;
334:       DMStagGetProductCoordinateArraysRead(dmstag,&cArrX,&cArrY,NULL);
335:       for (ey=start[1]; ey<start[1] + n[1] + extraPoint[1]; ++ey) {
336:         for (ex=start[0]; ex<start[0] + n[0] + extraPoint[0]; ++ex) {
337:           cArrDa[ey][ex][0] = cArrX[ex][0];
338:           cArrDa[ey][ex][1] = cArrY[ey][0];
339:         }
340:       }
341:       DMStagRestoreProductCoordinateArraysRead(dmstag,&cArrX,&cArrY,NULL);
342:     } else SETERRQ(PetscObjectComm((PetscObject)dmstag),PETSC_ERR_SUP,"Stag to DA coordinate transfer only supported for DMStag coordinate DM of type DMstag or DMProduct");
343:     DMDAVecRestoreArrayDOF(dmdaCoord,daCoord,&cArrDa);
344:   }  else if (dim == 3) {
345:     PetscInt ex,ey,ez;
346:     PetscScalar ****cArrDa;
347:     DMDAVecGetArrayDOF(dmdaCoord,daCoord,&cArrDa);
348:     if (daCoordIsStag)  {
349:       PetscInt slot;
350:       PetscScalar ****cArrStag;
351:       DMStagGetLocationSlot(dmstagCoord,loc,0,&slot);
352:       DMStagVecGetArrayRead(dmstagCoord,stagCoord,&cArrStag);
353:       for (ez=start[2]; ez<start[2] + n[2] + extraPoint[2]; ++ez) {
354:         for (ey=start[1]; ey<start[1] + n[1] + extraPoint[1]; ++ey) {
355:           for (ex=start[0]; ex<start[0] + n[0] + extraPoint[0]; ++ex) {
356:             for (d=0; d<3; ++d) {
357:               cArrDa[ez][ey][ex][d] = cArrStag[ez][ey][ex][slot+d];
358:             }
359:           }
360:         }
361:       }
362:       DMStagVecRestoreArrayRead(dmstagCoord,stagCoord,&cArrStag);
363:     } else if (daCoordIsProduct) {
364:       PetscScalar **cArrX,**cArrY,**cArrZ;
365:       DMStagGetProductCoordinateArraysRead(dmstag,&cArrX,&cArrY,&cArrZ);
366:       for (ez=start[2]; ez<start[2] + n[2] + extraPoint[2]; ++ez) {
367:         for (ey=start[1]; ey<start[1] + n[1] + extraPoint[1]; ++ey) {
368:           for (ex=start[0]; ex<start[0] + n[0] + extraPoint[0]; ++ex) {
369:             cArrDa[ez][ey][ex][0] = cArrX[ex][0];
370:             cArrDa[ez][ey][ex][1] = cArrY[ey][0];
371:             cArrDa[ez][ey][ex][2] = cArrZ[ez][0];
372:           }
373:         }
374:       }
375:       DMStagRestoreProductCoordinateArraysRead(dmstag,&cArrX,&cArrY,&cArrZ);
376:     } else SETERRQ(PetscObjectComm((PetscObject)dmstag),PETSC_ERR_SUP,"Stag to DA coordinate transfer only supported for DMStag coordinate DM of type DMstag or DMProduct");
377:     DMDAVecRestoreArrayDOF(dmdaCoord,daCoord,&cArrDa);
378:   } else SETERRQ(PetscObjectComm((PetscObject)dmstag),PETSC_ERR_SUP,"Unsupported dimension %d",dim);
379:   return 0;
380: }

382: /*@C
383:   DMStagVecSplitToDMDA - create a DMDA and Vec from a DMStag and Vec

385:   Logically Collective

387:   High-level helper function which accepts a DMStag, a global vector, and location/dof,
388:   and generates a corresponding DMDA and Vec.

390:   Input Parameters:
391: + dm - the DMStag object
392: . vec- Vec object associated with dm
393: . loc - which subgrid to extract (see DMStagStencilLocation)
394: - c - which component to extract (see note below)

396:   Output Parameters:
397: + pda - the new DMDA
398: - pdavec - the new Vec

400:   Notes:
401:   If a c value of -k is provided, the first k dof for that position are extracted,
402:   padding with zero values if needbe. If a non-negative value is provided, a single
403:   dof is extracted.

405:   The caller is responsible for destroying the created DMDA and Vec.

407:   Level: advanced

409: .seealso: DMSTAG, DMDA, DMStagMigrateVec(), DMStagCreateCompatibleDMStag()
410: @*/
411: PetscErrorCode DMStagVecSplitToDMDA(DM dm,Vec vec,DMStagStencilLocation loc,PetscInt c,DM *pda,Vec *pdavec)
412: {
413:   PetscInt        dim,locdof;
414:   DM              da,coordDM;
415:   Vec             davec;
416:   DMStagStencilLocation locCanonical;

420:   DMGetDimension(dm,&dim);
421:   DMStagGetLocationDOF(dm,loc,&locdof);
423:   DMStagStencilLocationCanonicalize(loc,&locCanonical);
424:   DMStagCreateCompatibleDMDA(dm,locCanonical,c,pda);
425:   da = *pda;
426:   DMSetUp(*pda);
427:   DMGetCoordinateDM(dm,&coordDM);
428:   if (coordDM) {
429:     DMStagTransferCoordinatesToDMDA(dm,locCanonical,da);
430:   }
431:   DMCreateGlobalVector(da,pdavec);
432:   davec = *pdavec;
433:   DMStagMigrateVecDMDA(dm,vec,locCanonical,c,da,davec);
434:   return 0;
435: }