Date: 2022-03-31T04:47:40+00:00

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