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

Actual source code: ex4.c

      1: static char help[] = "Solves the incompressible, variable viscosity Stokes equation in 2D driven by buouyancy variations.\n\n";

      3: #include <petscdm.h>
      4: #include <petscksp.h>
      5: #include <petscdmstag.h>
      6: #include <petscdmda.h>

      8: /* Shorter, more convenient names for DMStagStencilLocation entries */
      9: #define DOWN_LEFT  DMSTAG_DOWN_LEFT
     10: #define DOWN       DMSTAG_DOWN
     11: #define DOWN_RIGHT DMSTAG_DOWN_RIGHT
     12: #define LEFT       DMSTAG_LEFT
     13: #define ELEMENT    DMSTAG_ELEMENT
     14: #define RIGHT      DMSTAG_RIGHT
     15: #define UP_LEFT    DMSTAG_UP_LEFT
     16: #define UP         DMSTAG_UP
     17: #define UP_RIGHT   DMSTAG_UP_RIGHT

     19: /* An application context */
     20: typedef struct {
     21:   MPI_Comm    comm;
     22:   DM          dmStokes,dmCoeff;
     23:   Vec         coeff;
     24:   PetscReal   xmax,ymax,xmin,ymin,hxCharacteristic,hyCharacteristic;
     25:   PetscScalar eta1,eta2,rho1,rho2,gy,Kbound,Kcont,etaCharacteristic;
     26: } CtxData;
     27: typedef CtxData* Ctx;

     29: /* Helper functions */
     30: static PetscErrorCode PopulateCoefficientData(Ctx);
     31: static PetscErrorCode CreateSystem(const Ctx,Mat*,Vec*);
     32: static PetscErrorCode DumpSolution(Ctx,Vec);

     34: /* Coefficient/forcing Functions */
     35: static PetscScalar getRho(Ctx ctx,PetscScalar x) { return PetscRealPart(x) < (ctx->xmax-ctx->xmin)/2.0 ? ctx->rho1 : ctx->rho2; }
     36: static PetscScalar getEta(Ctx ctx,PetscScalar x) { return PetscRealPart(x) < (ctx->xmax-ctx->xmin)/2.0 ? ctx->eta1 : ctx->eta2; }

     38: int main(int argc,char **argv)
     39: {
     41:   Ctx            ctx;
     42:   Mat            A;
     43:   Vec            x,b;
     44:   KSP            ksp;

     46:   PetscInitialize(&argc,&argv,(char*)0,help);

     48:   /* Populate application context */
     49:   PetscMalloc1(1,&ctx);
     50:   ctx->comm = PETSC_COMM_WORLD;
     51:   ctx->xmin = 0.0;
     52:   ctx->xmax = 1e6;
     53:   ctx->ymin = 0.0;
     54:   ctx->ymax = 1.5e6;
     55:   ctx->rho1 = 3200;
     56:   ctx->rho2 = 3300;
     57:   ctx->eta1 = 1e20;
     58:   ctx->eta2 = 1e22;
     59:   ctx->gy   = 10.0;

     61:   /* Create two DMStag objects, corresponding to the same domain and parallel
     62:      decomposition ("topology"). Each defines a different set of fields on
     63:      the domain ("section"); the first the solution to the Stokes equations
     64:      (x- and y-velocities and scalar pressure), and the second holds coefficients
     65:      (viscosities on corners/elements and densities on corners) */
     66:   DMStagCreate2d(
     67:       ctx->comm,
     68:       DM_BOUNDARY_NONE,DM_BOUNDARY_NONE,
     69:       20,30,                                   /* Global element counts */
     70:       PETSC_DECIDE,PETSC_DECIDE,               /* Determine parallel decomposition automatically */
     71:       0,1,1,                                   /* dof: 0 per vertex, 1 per edge, 1 per face/element */
     72:       DMSTAG_STENCIL_BOX,
     73:       1,                                       /* elementwise stencil width */
     74:       NULL,NULL,
     75:       &ctx->dmStokes);
     76:   DMSetFromOptions(ctx->dmStokes);
     77:   DMSetUp(ctx->dmStokes);
     78:   DMStagSetUniformCoordinatesExplicit(ctx->dmStokes,0.0,ctx->xmax,0.0,ctx->ymax,0.0,0.0);
     79:   DMStagCreateCompatibleDMStag(ctx->dmStokes,2,0,1,0,&ctx->dmCoeff);
     80:   DMSetUp(ctx->dmCoeff);
     81:   DMStagSetUniformCoordinatesExplicit(ctx->dmCoeff,0.0,ctx->xmax,0.0,ctx->ymax,0.0,0.0);

     83:   /* Note: see ex2.c for a more-efficient way to work with coordinates on an
     84:      orthogonal grid, using DMStagSetUniformCoordinatesProduct() */

     86:   /* Get scaling constants, knowing grid spacing */
     87:   {
     88:     PetscInt N[2];
     89:     PetscReal hxAvgInv;
     90:     DMStagGetGlobalSizes(ctx->dmStokes,&N[0],&N[1],NULL);
     91:     ctx->hxCharacteristic = (ctx->xmax-ctx->xmin)/N[0];
     92:     ctx->hyCharacteristic = (ctx->ymax-ctx->ymin)/N[1];
     93:     ctx->etaCharacteristic = PetscMin(PetscRealPart(ctx->eta1),PetscRealPart(ctx->eta2));
     94:     hxAvgInv = 2.0/(ctx->hxCharacteristic + ctx->hyCharacteristic);
     95:     ctx->Kcont = ctx->etaCharacteristic*hxAvgInv;
     96:     ctx->Kbound = ctx->etaCharacteristic*hxAvgInv*hxAvgInv;
     97:   }

     99:   /* Populate coefficient data */
    100:   PopulateCoefficientData(ctx);

    102:   /* Construct System */
    103:   CreateSystem(ctx,&A,&b);

    105:   /* Solve */
    106:   VecDuplicate(b,&x);
    107:   KSPCreate(ctx->comm,&ksp);
    108:   KSPSetType(ksp,KSPFGMRES);
    109:   KSPSetOperators(ksp,A,A);
    110:   PetscOptionsSetValue(NULL,"-ksp_converged_reason",""); /* To get info on direct solve success */
    111:   KSPSetDM(ksp,ctx->dmStokes);
    112:   KSPSetDMActive(ksp,PETSC_FALSE);
    113:   KSPSetFromOptions(ksp);
    114:   KSPSolve(ksp,b,x);
    115:   {
    116:     KSPConvergedReason reason;
    117:     KSPGetConvergedReason(ksp,&reason);
    119:   }

    121:   /* Dump solution by converting to DMDAs and dumping */
    122:   DumpSolution(ctx,x);

    124:   /* Destroy PETSc objects and finalize */
    125:   MatDestroy(&A);
    126:   VecDestroy(&x);
    127:   VecDestroy(&b);
    128:   VecDestroy(&ctx->coeff);
    129:   KSPDestroy(&ksp);
    130:   DMDestroy(&ctx->dmStokes);
    131:   DMDestroy(&ctx->dmCoeff);
    132:   PetscFree(ctx);
    133:   PetscFinalize();
    134:   return 0;
    135: }

    137: static PetscErrorCode CreateSystem(const Ctx ctx,Mat *pA,Vec *pRhs)
    138: {
    139:   PetscInt       N[2];
    140:   PetscInt       ex,ey,startx,starty,nx,ny;
    141:   Mat            A;
    142:   Vec            rhs;
    143:   PetscReal      hx,hy;
    144:   const PetscBool pinPressure = PETSC_TRUE;
    145:   Vec            coeffLocal;

    148:   DMCreateMatrix(ctx->dmStokes,pA);
    149:   A = *pA;
    150:   DMCreateGlobalVector(ctx->dmStokes,pRhs);
    151:   rhs = *pRhs;
    152:   DMStagGetCorners(ctx->dmStokes,&startx,&starty,NULL,&nx,&ny,NULL,NULL,NULL,NULL);
    153:   DMStagGetGlobalSizes(ctx->dmStokes,&N[0],&N[1],NULL);
    154:   hx = ctx->hxCharacteristic;
    155:   hy = ctx->hyCharacteristic;
    156:   DMGetLocalVector(ctx->dmCoeff,&coeffLocal);
    157:   DMGlobalToLocal(ctx->dmCoeff,ctx->coeff,INSERT_VALUES,coeffLocal);

    159:   /* Loop over all local elements. Note that it may be more efficient in real
    160:      applications to loop over each boundary separately */
    161:   for (ey = starty; ey<starty+ny; ++ey) { /* With DMStag, always iterate x fastest, y second fastest, z slowest */
    162:     for (ex = startx; ex<startx+nx; ++ex) {

    164:       if (ey == N[1]-1) {
    165:         /* Top boundary velocity Dirichlet */
    166:         DMStagStencil row;
    167:         PetscScalar   valRhs;
    168:         const PetscScalar valA = ctx->Kbound;
    169:         row.i = ex; row.j = ey; row.loc = UP; row.c = 0;
    170:         DMStagMatSetValuesStencil(ctx->dmStokes,A,1,&row,1,&row,&valA,INSERT_VALUES);
    171:         valRhs = 0.0;
    172:         DMStagVecSetValuesStencil(ctx->dmStokes,rhs,1,&row,&valRhs,INSERT_VALUES);
    173:       }

    175:       if (ey == 0) {
    176:         /* Bottom boundary velocity Dirichlet */
    177:         DMStagStencil row;
    178:         PetscScalar   valRhs;
    179:         const PetscScalar valA = ctx->Kbound;
    180:         row.i = ex; row.j = ey; row.loc = DOWN; row.c = 0;
    181:         DMStagMatSetValuesStencil(ctx->dmStokes,A,1,&row,1,&row,&valA,INSERT_VALUES);
    182:         valRhs = 0.0;
    183:         DMStagVecSetValuesStencil(ctx->dmStokes,rhs,1,&row,&valRhs,INSERT_VALUES);
    184:       } else {
    185:         /* Y-momentum equation : (u_xx + u_yy) - p_y = f^y : includes non-zero forcing */
    186:         PetscInt      nEntries;
    187:         DMStagStencil row,col[11];
    188:         PetscScalar   valA[11];
    189:         DMStagStencil rhoPoint[2];
    190:         PetscScalar   rho[2],valRhs;
    191:         DMStagStencil etaPoint[4];
    192:         PetscScalar   eta[4],etaLeft,etaRight,etaUp,etaDown;

    194:         /* get rho values  and compute rhs value*/
    195:         rhoPoint[0].i = ex; rhoPoint[0].j = ey; rhoPoint[0].loc = DOWN_LEFT;  rhoPoint[0].c = 1;
    196:         rhoPoint[1].i = ex; rhoPoint[1].j = ey; rhoPoint[1].loc = DOWN_RIGHT; rhoPoint[1].c = 1;
    197:         DMStagVecGetValuesStencil(ctx->dmCoeff,coeffLocal,2,rhoPoint,rho);
    198:         valRhs = -ctx->gy * 0.5 * (rho[0] + rho[1]);

    200:         /* Get eta values */
    201:         etaPoint[0].i = ex; etaPoint[0].j = ey;   etaPoint[0].loc = DOWN_LEFT;  etaPoint[0].c = 0; /* Left  */
    202:         etaPoint[1].i = ex; etaPoint[1].j = ey;   etaPoint[1].loc = DOWN_RIGHT; etaPoint[1].c = 0; /* Right */
    203:         etaPoint[2].i = ex; etaPoint[2].j = ey;   etaPoint[2].loc = ELEMENT;    etaPoint[2].c = 0; /* Up    */
    204:         etaPoint[3].i = ex; etaPoint[3].j = ey-1; etaPoint[3].loc = ELEMENT;    etaPoint[3].c = 0; /* Down  */
    205:         DMStagVecGetValuesStencil(ctx->dmCoeff,coeffLocal,4,etaPoint,eta);
    206:         etaLeft = eta[0]; etaRight = eta[1]; etaUp = eta[2]; etaDown = eta[3];

    208:         if (ex == 0) {
    209:           /* Left boundary y velocity stencil */
    210:           nEntries = 10;
    211:           row.i    = ex  ; row.j    = ey  ; row.loc    = DOWN;     row.c     = 0;
    212:           col[0].i = ex  ; col[0].j = ey  ; col[0].loc = DOWN;     col[0].c  = 0; valA[0]  = -2.0 * (etaDown + etaUp) / (hy*hy) - (etaRight) /(hx*hx);
    213:           col[1].i = ex  ; col[1].j = ey-1; col[1].loc = DOWN;     col[1].c  = 0; valA[1]  =  2.0 * etaDown  / (hy*hy);
    214:           col[2].i = ex  ; col[2].j = ey+1; col[2].loc = DOWN;     col[2].c  = 0; valA[2]  =  2.0 * etaUp    / (hy*hy);
    215:           /* No left entry */
    216:           col[3].i = ex+1; col[3].j = ey  ; col[3].loc = DOWN;     col[3].c  = 0; valA[3]  =        etaRight / (hx*hx);
    217:           col[4].i = ex  ; col[4].j = ey-1; col[4].loc = LEFT;     col[4].c  = 0; valA[4]  =        etaLeft  / (hx*hy); /* down left x edge */
    218:           col[5].i = ex  ; col[5].j = ey-1; col[5].loc = RIGHT;    col[5].c  = 0; valA[5]  = -      etaRight / (hx*hy); /* down right x edge */
    219:           col[6].i = ex  ; col[6].j = ey  ; col[6].loc = LEFT;     col[6].c  = 0; valA[6]  = -      etaLeft  / (hx*hy); /* up left x edge */
    220:           col[7].i = ex  ; col[7].j = ey  ; col[7].loc = RIGHT;    col[7].c  = 0; valA[7]  =        etaRight / (hx*hy); /* up right x edge */
    221:           col[8].i = ex  ; col[8].j = ey-1; col[8].loc = ELEMENT;  col[8].c  = 0; valA[8]  =  ctx->Kcont / hy;
    222:           col[9].i = ex  ; col[9].j = ey  ; col[9].loc = ELEMENT;  col[9].c  = 0; valA[9]  = -ctx->Kcont / hy;
    223:         } else if (ex == N[0]-1) {
    224:           /* Right boundary y velocity stencil */
    225:           nEntries = 10;
    226:           row.i    = ex  ; row.j    = ey  ; row.loc    = DOWN;     row.c     = 0;
    227:           col[0].i = ex  ; col[0].j = ey  ; col[0].loc = DOWN;     col[0].c  = 0; valA[0]  = -2.0 * (etaDown + etaUp) / (hy*hy) - (etaLeft) /(hx*hx);
    228:           col[1].i = ex  ; col[1].j = ey-1; col[1].loc = DOWN;     col[1].c  = 0; valA[1]  =  2.0 * etaDown  / (hy*hy);
    229:           col[2].i = ex  ; col[2].j = ey+1; col[2].loc = DOWN;     col[2].c  = 0; valA[2]  =  2.0 * etaUp    / (hy*hy);
    230:           col[3].i = ex-1; col[3].j = ey  ; col[3].loc = DOWN;     col[3].c  = 0; valA[3]  =        etaLeft  / (hx*hx);
    231:           /* No right element */
    232:           col[4].i = ex  ; col[4].j = ey-1; col[4].loc = LEFT;     col[4].c  = 0; valA[4]  =        etaLeft  / (hx*hy); /* down left x edge */
    233:           col[5].i = ex  ; col[5].j = ey-1; col[5].loc = RIGHT;    col[5].c  = 0; valA[5]  = -      etaRight / (hx*hy); /* down right x edge */
    234:           col[6].i = ex  ; col[6].j = ey  ; col[6].loc = LEFT;     col[6].c  = 0; valA[7]  = -      etaLeft  / (hx*hy); /* up left x edge */
    235:           col[7].i = ex  ; col[7].j = ey  ; col[7].loc = RIGHT;    col[7].c  = 0; valA[7]  =        etaRight / (hx*hy); /* up right x edge */
    236:           col[8].i = ex  ; col[8].j = ey-1; col[8].loc = ELEMENT;  col[8].c  = 0; valA[8]  =  ctx->Kcont / hy;
    237:           col[9].i = ex  ; col[9].j = ey  ; col[9].loc = ELEMENT;  col[9].c  = 0; valA[9]  = -ctx->Kcont / hy;
    238:         } else {
    239:           /* U_y interior equation */
    240:           nEntries = 11;
    241:           row.i    = ex  ; row.j     = ey  ; row.loc     = DOWN;     row.c     = 0;
    242:           col[0].i = ex  ; col[0].j  = ey  ; col[0].loc  = DOWN;     col[0].c  = 0; valA[0]  = -2.0 * (etaDown + etaUp) / (hy*hy) - (etaLeft + etaRight) /(hx*hx);
    243:           col[1].i = ex  ; col[1].j  = ey-1; col[1].loc  = DOWN;     col[1].c  = 0; valA[1]  =  2.0 * etaDown  / (hy*hy);
    244:           col[2].i = ex  ; col[2].j  = ey+1; col[2].loc  = DOWN;     col[2].c  = 0; valA[2]  =  2.0 * etaUp    / (hy*hy);
    245:           col[3].i = ex-1; col[3].j  = ey  ; col[3].loc  = DOWN;     col[3].c  = 0; valA[3]  =        etaLeft  / (hx*hx);
    246:           col[4].i = ex+1; col[4].j  = ey  ; col[4].loc  = DOWN;     col[4].c  = 0; valA[4]  =        etaRight / (hx*hx);
    247:           col[5].i = ex  ; col[5].j  = ey-1; col[5].loc  = LEFT;     col[5].c  = 0; valA[5]  =        etaLeft  / (hx*hy); /* down left x edge */
    248:           col[6].i = ex  ; col[6].j  = ey-1; col[6].loc  = RIGHT;    col[6].c  = 0; valA[6]  = -      etaRight / (hx*hy); /* down right x edge */
    249:           col[7].i = ex  ; col[7].j  = ey  ; col[7].loc  = LEFT;     col[7].c  = 0; valA[7]  = -      etaLeft  / (hx*hy); /* up left x edge */
    250:           col[8].i = ex  ; col[8].j  = ey  ; col[8].loc  = RIGHT;    col[8].c  = 0; valA[8]  =        etaRight / (hx*hy); /* up right x edge */
    251:           col[9].i = ex  ; col[9].j  = ey-1; col[9].loc  = ELEMENT;  col[9].c  = 0; valA[9]  =  ctx->Kcont / hy;
    252:           col[10].i = ex ; col[10].j = ey  ; col[10].loc = ELEMENT; col[10].c  = 0; valA[10] = -ctx->Kcont / hy;
    253:         }

    255:         /* Insert Y-momentum entries */
    256:         DMStagMatSetValuesStencil(ctx->dmStokes,A,1,&row,nEntries,col,valA,INSERT_VALUES);
    257:         DMStagVecSetValuesStencil(ctx->dmStokes,rhs,1,&row,&valRhs,INSERT_VALUES);
    258:       }

    260:       if (ex == N[0]-1) {
    261:         /* Right Boundary velocity Dirichlet */
    262:         /* Redundant in the corner */
    263:         DMStagStencil row;
    264:         PetscScalar   valRhs;

    266:         const PetscScalar valA = ctx->Kbound;
    267:         row.i = ex; row.j = ey; row.loc = RIGHT; row.c = 0;
    268:         DMStagMatSetValuesStencil(ctx->dmStokes,A,1,&row,1,&row,&valA,INSERT_VALUES);
    269:         valRhs = 0.0;
    270:         DMStagVecSetValuesStencil(ctx->dmStokes,rhs,1,&row,&valRhs,INSERT_VALUES);
    271:       }
    272:       if (ex == 0) {
    273:         /* Left velocity Dirichlet */
    274:         DMStagStencil row;
    275:         PetscScalar   valRhs;
    276:         const PetscScalar valA = ctx->Kbound;
    277:         row.i = ex; row.j = ey; row.loc = LEFT; row.c = 0;
    278:         DMStagMatSetValuesStencil(ctx->dmStokes,A,1,&row,1,&row,&valA,INSERT_VALUES);
    279:         valRhs = 0.0;
    280:         DMStagVecSetValuesStencil(ctx->dmStokes,rhs,1,&row,&valRhs,INSERT_VALUES);
    281:       } else {
    282:         /* X-momentum equation : (u_xx + u_yy) - p_x = f^x */
    283:         PetscInt nEntries;
    284:         DMStagStencil row,col[11];
    285:         PetscScalar   valRhs,valA[11];
    286:         DMStagStencil etaPoint[4];
    287:         PetscScalar eta[4],etaLeft,etaRight,etaUp,etaDown;

    289:         /* Get eta values */
    290:         etaPoint[0].i = ex-1; etaPoint[0].j = ey; etaPoint[0].loc = ELEMENT;   etaPoint[0].c = 0; /* Left  */
    291:         etaPoint[1].i = ex;   etaPoint[1].j = ey; etaPoint[1].loc = ELEMENT;   etaPoint[1].c = 0; /* Right */
    292:         etaPoint[2].i = ex;   etaPoint[2].j = ey; etaPoint[2].loc = UP_LEFT;   etaPoint[2].c = 0; /* Up    */
    293:         etaPoint[3].i = ex;   etaPoint[3].j = ey; etaPoint[3].loc = DOWN_LEFT; etaPoint[3].c = 0; /* Down  */
    294:         DMStagVecGetValuesStencil(ctx->dmCoeff,coeffLocal,4,etaPoint,eta);
    295:         etaLeft = eta[0]; etaRight = eta[1]; etaUp = eta[2]; etaDown = eta[3];

    297:         if (ey == 0) {
    298:           /* Bottom boundary x velocity stencil (with zero vel deriv) */
    299:           nEntries = 10;
    300:           row.i     = ex  ; row.j     = ey  ; row.loc     = LEFT;    row.c      = 0;
    301:           col[0].i  = ex  ; col[0].j  = ey  ; col[0].loc  = LEFT;    col[0].c   = 0; valA[0]  = -2.0 * (etaLeft + etaRight) / (hx*hx) -(etaUp) / (hy*hy);
    302:           /* Missing element below */
    303:           col[1].i = ex  ; col[1].j  = ey+1; col[1].loc  = LEFT;    col[1].c   = 0; valA[1]  =        etaUp    / (hy*hy);
    304:           col[2].i = ex-1; col[2].j  = ey  ; col[2].loc  = LEFT;    col[2].c   = 0; valA[2]  =  2.0 * etaLeft  / (hx*hx);
    305:           col[3].i = ex+1; col[3].j  = ey  ; col[3].loc  = LEFT;    col[3].c   = 0; valA[3]  =  2.0 * etaRight / (hx*hx);
    306:           col[4].i = ex-1; col[4].j  = ey  ; col[4].loc  = DOWN;    col[4].c   = 0; valA[4]  =        etaDown  / (hx*hy); /* down left */
    307:           col[5].i = ex  ; col[5].j  = ey  ; col[5].loc  = DOWN;    col[5].c   = 0; valA[5]  = -      etaDown  / (hx*hy); /* down right */
    308:           col[6].i = ex-1; col[6].j  = ey  ; col[6].loc  = UP;      col[6].c   = 0; valA[6]  = -      etaUp    / (hx*hy); /* up left */
    309:           col[7].i = ex  ; col[7].j  = ey  ; col[7].loc  = UP;      col[7].c   = 0; valA[7]  =        etaUp    / (hx*hy); /* up right */
    310:           col[8].i = ex-1; col[8].j  = ey  ; col[8].loc  = ELEMENT; col[8].c   = 0; valA[8]  =  ctx->Kcont / hx;
    311:           col[9].i = ex  ; col[9].j  = ey  ; col[9].loc  = ELEMENT; col[9].c   = 0; valA[9]  = -ctx->Kcont / hx;
    312:           valRhs = 0.0;
    313:         } else if (ey == N[1]-1) {
    314:           /* Top boundary x velocity stencil */
    315:           nEntries = 10;
    316:           row.i    = ex  ; row.j     = ey  ; row.loc     = LEFT;    row.c     = 0;
    317:           col[0].i = ex  ; col[0].j  = ey  ; col[0].loc  = LEFT;    col[0].c  = 0; valA[0]  = -2.0 * (etaLeft + etaRight) / (hx*hx) -(etaDown) / (hy*hy);
    318:           col[1].i = ex  ; col[1].j  = ey-1; col[1].loc  = LEFT;    col[1].c  = 0; valA[1]  =        etaDown  / (hy*hy);
    319:           /* Missing element above */
    320:           col[2].i = ex-1; col[2].j  = ey  ; col[2].loc  = LEFT;    col[2].c  = 0; valA[2]  =  2.0 * etaLeft  / (hx*hx);
    321:           col[3].i = ex+1; col[3].j  = ey  ; col[3].loc  = LEFT;    col[3].c  = 0; valA[3]  =  2.0 * etaRight / (hx*hx);
    322:           col[4].i = ex-1; col[4].j  = ey  ; col[4].loc  = DOWN;    col[4].c  = 0; valA[4]  =        etaDown  / (hx*hy); /* down left */
    323:           col[5].i = ex  ; col[5].j  = ey  ; col[5].loc  = DOWN;    col[5].c  = 0; valA[5]  = -      etaDown  / (hx*hy); /* down right */
    324:           col[6].i = ex-1; col[6].j  = ey  ; col[6].loc  = UP;      col[6].c  = 0; valA[6]  = -      etaUp    / (hx*hy); /* up left */
    325:           col[7].i = ex  ; col[7].j  = ey  ; col[7].loc  = UP;      col[7].c  = 0; valA[7]  =        etaUp    / (hx*hy); /* up right */
    326:           col[8].i = ex-1; col[8].j  = ey  ; col[8].loc  = ELEMENT; col[8].c  = 0; valA[8]  =  ctx->Kcont / hx;
    327:           col[9].i = ex  ; col[9].j  = ey  ; col[9].loc = ELEMENT;  col[9].c  = 0; valA[9]  = -ctx->Kcont / hx;
    328:           valRhs = 0.0;
    329:         } else {
    330:           /* U_x interior equation */
    331:           nEntries = 11;
    332:           row.i     = ex  ; row.j     = ey  ; row.loc     = LEFT;    row.c      = 0;
    333:           col[0].i  = ex  ; col[0].j  = ey  ; col[0].loc  = LEFT;    col[0].c   = 0; valA[0]  = -2.0 * (etaLeft + etaRight) / (hx*hx) -(etaUp + etaDown) / (hy*hy);
    334:           col[1].i  = ex  ; col[1].j  = ey-1; col[1].loc  = LEFT;    col[1].c   = 0; valA[1]  =        etaDown  / (hy*hy);
    335:           col[2].i  = ex  ; col[2].j  = ey+1; col[2].loc  = LEFT;    col[2].c   = 0; valA[2]  =        etaUp    / (hy*hy);
    336:           col[3].i  = ex-1; col[3].j  = ey  ; col[3].loc  = LEFT;    col[3].c   = 0; valA[3]  =  2.0 * etaLeft  / (hx*hx);
    337:           col[4].i  = ex+1; col[4].j  = ey  ; col[4].loc  = LEFT;    col[4].c   = 0; valA[4]  =  2.0 * etaRight / (hx*hx);
    338:           col[5].i  = ex-1; col[5].j  = ey  ; col[5].loc  = DOWN;    col[5].c   = 0; valA[5]  =        etaDown  / (hx*hy); /* down left */
    339:           col[6].i  = ex  ; col[6].j  = ey  ; col[6].loc  = DOWN;    col[6].c   = 0; valA[6]  = -      etaDown  / (hx*hy); /* down right */
    340:           col[7].i  = ex-1; col[7].j  = ey  ; col[7].loc  = UP;      col[7].c   = 0; valA[7]  = -      etaUp    / (hx*hy); /* up left */
    341:           col[8].i  = ex  ; col[8].j  = ey  ; col[8].loc  = UP;      col[8].c   = 0; valA[8]  =        etaUp    / (hx*hy); /* up right */
    342:           col[9].i  = ex-1; col[9].j  = ey  ; col[9].loc  = ELEMENT; col[9].c   = 0; valA[9]  =  ctx->Kcont / hx;
    343:           col[10].i = ex  ; col[10].j = ey  ; col[10].loc = ELEMENT; col[10].c  = 0; valA[10] = -ctx->Kcont / hx;
    344:           valRhs = 0.0;
    345:         }
    346:         DMStagMatSetValuesStencil(ctx->dmStokes,A,1,&row,nEntries,col,valA,INSERT_VALUES);
    347:         DMStagVecSetValuesStencil(ctx->dmStokes,rhs,1,&row,&valRhs,INSERT_VALUES);
    348:       }

    350:       /* P equation : u_x + v_y = 0
    351:          Note that this includes an explicit zero on the diagonal. This is only needed for
    352:          direct solvers (not required if using an iterative solver and setting the constant-pressure nullspace) */
    353:       if (pinPressure && ex == 0 && ey == 0) { /* Pin the first pressure node to zero, if requested */
    354:         DMStagStencil row;
    355:         PetscScalar valA,valRhs;
    356:         row.i = ex; row.j = ey; row.loc = ELEMENT; row.c = 0;
    357:         valA = ctx->Kbound;
    358:         DMStagMatSetValuesStencil(ctx->dmStokes,A,1,&row,1,&row,&valA,INSERT_VALUES);
    359:         valRhs = 0.0;
    360:         DMStagVecSetValuesStencil(ctx->dmStokes,rhs,1,&row,&valRhs,INSERT_VALUES);
    361:       } else {
    362:         DMStagStencil row,col[5];
    363:         PetscScalar   valA[5],valRhs;

    365:         row.i    = ex; row.j    = ey; row.loc    = ELEMENT; row.c    = 0;
    366:         col[0].i = ex; col[0].j = ey; col[0].loc = LEFT;    col[0].c = 0; valA[0] = -ctx->Kcont / hx;
    367:         col[1].i = ex; col[1].j = ey; col[1].loc = RIGHT;   col[1].c = 0; valA[1] =  ctx->Kcont / hx;
    368:         col[2].i = ex; col[2].j = ey; col[2].loc = DOWN;    col[2].c = 0; valA[2] = -ctx->Kcont / hy;
    369:         col[3].i = ex; col[3].j = ey; col[3].loc = UP;      col[3].c = 0; valA[3] =  ctx->Kcont / hy;
    370:         col[4] = row;                                                     valA[4] = 0.0;
    371:         DMStagMatSetValuesStencil(ctx->dmStokes,A,1,&row,5,col,valA,INSERT_VALUES);
    372:         valRhs = 0.0;
    373:         DMStagVecSetValuesStencil(ctx->dmStokes,rhs,1,&row,&valRhs,INSERT_VALUES);
    374:       }
    375:     }
    376:   }
    377:   DMRestoreLocalVector(ctx->dmCoeff,&coeffLocal);
    378:   MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);
    379:   VecAssemblyBegin(rhs);
    380:   MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);
    381:   VecAssemblyEnd(rhs);
    382:   return 0;
    383: }

    385: /* Here, we demonstrate getting coordinates from a vector by using DMStagStencil.
    386: This would usually be done with direct array access, though. */
    387: static PetscErrorCode PopulateCoefficientData(Ctx ctx)
    388: {
    389:   PetscInt       N[2],nExtra[2];
    390:   PetscInt       ex,ey,startx,starty,nx,ny;
    391:   Vec            coeffLocal,coordLocal;
    392:   DM             dmCoord;

    395:   DMCreateGlobalVector(ctx->dmCoeff,&ctx->coeff);
    396:   DMGetLocalVector(ctx->dmCoeff,&coeffLocal);
    397:   DMStagGetCorners(ctx->dmCoeff,&startx,&starty,NULL,&nx,&ny,NULL,&nExtra[0],&nExtra[1],NULL);
    398:   DMStagGetGlobalSizes(ctx->dmCoeff,&N[0],&N[1],NULL);
    399:   DMGetCoordinatesLocal(ctx->dmCoeff,&coordLocal);
    400:   DMGetCoordinateDM(ctx->dmCoeff,&dmCoord);
    401:   for (ey = starty; ey<starty+ny+nExtra[1]; ++ey) {
    402:     for (ex = startx; ex<startx+nx+nExtra[0]; ++ex) {

    404:       /* Eta (element) */
    405:       if (ey < starty + ny && ex < startx + nx) {
    406:         DMStagStencil point,pointCoordx;
    407:         PetscScalar   val,x;
    408:         point.i = ex; point.j = ey; point.loc = ELEMENT; point.c = 0;
    409:         pointCoordx = point;
    410:         DMStagVecGetValuesStencil(dmCoord,coordLocal,1,&pointCoordx,&x);
    411:         val = getEta(ctx,x);
    412:         DMStagVecSetValuesStencil(ctx->dmCoeff,ctx->coeff,1,&point,&val,INSERT_VALUES);
    413:       }

    415:       /* Rho */
    416:       {
    417:         DMStagStencil point,pointCoordx;
    418:         PetscScalar   val,x;
    419:         point.i = ex; point.j = ey; point.loc = DOWN_LEFT; point.c = 1; /* Note .c = 1 */
    420:         pointCoordx = point; pointCoordx.c = 0;
    421:         DMStagVecGetValuesStencil(dmCoord,coordLocal,1,&pointCoordx,&x);
    422:         val = getRho(ctx,x);
    423:         DMStagVecSetValuesStencil(ctx->dmCoeff,ctx->coeff,1,&point,&val,INSERT_VALUES);
    424:       }

    426:       /* Eta (corner) - populate extra corners on right/top of domain */
    427:       {
    428:         DMStagStencil point,pointCoordx;
    429:         PetscScalar   val,x;
    430:         point.i = ex; point.j = ey; point.loc = DOWN_LEFT; point.c = 0;
    431:         pointCoordx = point;
    432:         DMStagVecGetValuesStencil(dmCoord,coordLocal,1,&pointCoordx,&x);
    433:         val = getEta(ctx,x);
    434:         DMStagVecSetValuesStencil(ctx->dmCoeff,ctx->coeff,1,&point,&val,INSERT_VALUES);
    435:       }
    436:       if (ex == N[0]-1) {
    437:         DMStagStencil point,pointCoordx;
    438:         PetscScalar   val,x;
    439:         point.i = ex; point.j = ey; point.loc = DOWN_RIGHT; point.c = 0;
    440:         pointCoordx = point;
    441:         DMStagVecGetValuesStencil(dmCoord,coordLocal,1,&pointCoordx,&x);
    442:         val = getEta(ctx,x);
    443:         DMStagVecSetValuesStencil(ctx->dmCoeff,ctx->coeff,1,&point,&val,INSERT_VALUES);
    444:       }
    445:       if (ey == N[1]-1) {
    446:         DMStagStencil point,pointCoordx;
    447:         PetscScalar   val,x;
    448:         point.i = ex; point.j = ey; point.loc = UP_LEFT; point.c = 0;
    449:         pointCoordx = point;
    450:         DMStagVecGetValuesStencil(dmCoord,coordLocal,1,&pointCoordx,&x);
    451:         val = getEta(ctx,x);
    452:         DMStagVecSetValuesStencil(ctx->dmCoeff,ctx->coeff,1,&point,&val,INSERT_VALUES);
    453:       }
    454:       if (ex == N[0]-1 && ey == N[1]-1) {
    455:         DMStagStencil point,pointCoordx;
    456:         PetscScalar   val,x;
    457:         point.i = ex; point.j = ey; point.loc = UP_RIGHT; point.c = 0;
    458:         pointCoordx = point;
    459:         DMStagVecGetValuesStencil(dmCoord,coordLocal,1,&pointCoordx,&x);
    460:         val = getEta(ctx,x);
    461:         DMStagVecSetValuesStencil(ctx->dmCoeff,ctx->coeff,1,&point,&val,INSERT_VALUES);
    462:       }
    463:     }
    464:   }
    465:   VecAssemblyBegin(ctx->coeff);
    466:   VecAssemblyEnd(ctx->coeff);
    467:   DMRestoreLocalVector(ctx->dmCoeff,&coeffLocal);
    468:   return 0;
    469: }

    471: static PetscErrorCode DumpSolution(Ctx ctx,Vec x)
    472: {
    473:   DM             dmVelAvg;
    474:   Vec            velAvg;
    475:   DM             daVelAvg,daP,daEtaElement,daEtaCorner,daRho;
    476:   Vec            vecVelAvg,vecP,vecEtaElement,vecEtaCorner,vecRho;

    480:   /* For convenience, create a new DM and Vec which will hold averaged velocities
    481:      Note that this could also be accomplished with direct array access, using
    482:      DMStagVecGetArray() and related functions */
    483:   DMStagCreateCompatibleDMStag(ctx->dmStokes,0,0,2,0,&dmVelAvg); /* 2 dof per element */
    484:   DMSetUp(dmVelAvg);
    485:   DMStagSetUniformCoordinatesExplicit(dmVelAvg,0.0,ctx->xmax,0.0,ctx->ymax,0.0,0.0);
    486:   DMCreateGlobalVector(dmVelAvg,&velAvg);
    487:   {
    488:     PetscInt ex,ey,startx,starty,nx,ny;
    489:     Vec      stokesLocal;
    490:     DMGetLocalVector(ctx->dmStokes,&stokesLocal);
    491:     DMGlobalToLocal(ctx->dmStokes,x,INSERT_VALUES,stokesLocal);
    492:     DMStagGetCorners(dmVelAvg,&startx,&starty,NULL,&nx,&ny,NULL,NULL,NULL,NULL);
    493:     for (ey = starty; ey<starty+ny; ++ey) {
    494:       for (ex = startx; ex<startx+nx; ++ex) {
    495:         DMStagStencil from[4],to[2];
    496:         PetscScalar   valFrom[4],valTo[2];
    497:         from[0].i = ex; from[0].j = ey; from[0].loc = UP;    from[0].c = 0;
    498:         from[1].i = ex; from[1].j = ey; from[1].loc = DOWN;  from[1].c = 0;
    499:         from[2].i = ex; from[2].j = ey; from[2].loc = LEFT;  from[2].c = 0;
    500:         from[3].i = ex; from[3].j = ey; from[3].loc = RIGHT; from[3].c = 0;
    501:         DMStagVecGetValuesStencil(ctx->dmStokes,stokesLocal,4,from,valFrom);
    502:         to[0].i = ex; to[0].j = ey; to[0].loc = ELEMENT;    to[0].c = 0; valTo[0] = 0.5 * (valFrom[2] + valFrom[3]);
    503:         to[1].i = ex; to[1].j = ey; to[1].loc = ELEMENT;    to[1].c = 1; valTo[1] = 0.5 * (valFrom[0] + valFrom[1]);
    504:         DMStagVecSetValuesStencil(dmVelAvg,velAvg,2,to,valTo,INSERT_VALUES);
    505:       }
    506:     }
    507:     VecAssemblyBegin(velAvg);
    508:     VecAssemblyEnd(velAvg);
    509:     DMRestoreLocalVector(ctx->dmStokes,&stokesLocal);
    510:   }

    512:   /* Create individual DMDAs for sub-grids of our DMStag objects. This is
    513:      somewhat inefficient, but allows use of the DMDA API without re-implementing
    514:      all utilities for DMStag */

    516:   DMStagVecSplitToDMDA(ctx->dmStokes,x,DMSTAG_ELEMENT,0,&daP,&vecP);
    517:   PetscObjectSetName((PetscObject)vecP,"p (scaled)");
    518:   DMStagVecSplitToDMDA(ctx->dmCoeff,ctx->coeff,DMSTAG_DOWN_LEFT,0, &daEtaCorner, &vecEtaCorner);
    519:   PetscObjectSetName((PetscObject)vecEtaCorner,"eta");
    520:   DMStagVecSplitToDMDA(ctx->dmCoeff,ctx->coeff,DMSTAG_ELEMENT,  0, &daEtaElement,&vecEtaElement);
    521:   PetscObjectSetName((PetscObject)vecEtaElement,"eta");
    522:   DMStagVecSplitToDMDA(ctx->dmCoeff,ctx->coeff,DMSTAG_DOWN_LEFT,  1, &daRho,       &vecRho);
    523:   PetscObjectSetName((PetscObject)vecRho,"density");
    524:   DMStagVecSplitToDMDA(dmVelAvg,    velAvg,    DMSTAG_ELEMENT,  -3,&daVelAvg,    &vecVelAvg); /* note -3 : pad with zero */
    525:   PetscObjectSetName((PetscObject)vecVelAvg,"Velocity (Averaged)");

    527:   /* Dump element-based fields to a .vtr file */
    528:   {
    529:     PetscViewer viewer;
    530:     PetscViewerVTKOpen(PetscObjectComm((PetscObject)daVelAvg),"ex4_element.vtr",FILE_MODE_WRITE,&viewer);
    531:     VecView(vecVelAvg,viewer);
    532:     VecView(vecP,viewer);
    533:     VecView(vecEtaElement,viewer);
    534:     PetscViewerDestroy(&viewer);
    535:   }

    537:   /* Dump vertex-based fields to a second .vtr file */
    538:   {
    539:     PetscViewer viewer;
    540:     PetscViewerVTKOpen(PetscObjectComm((PetscObject)daEtaCorner),"ex4_vertex.vtr",FILE_MODE_WRITE,&viewer);
    541:     VecView(vecEtaCorner,viewer);
    542:     VecView(vecRho,viewer);
    543:     PetscViewerDestroy(&viewer);
    544:   }

    546:   /* Edge-based fields could similarly be dumped */

    548:   /* Destroy DMDAs and Vecs */
    549:   VecDestroy(&vecVelAvg);
    550:   VecDestroy(&vecP);
    551:   VecDestroy(&vecEtaCorner);
    552:   VecDestroy(&vecEtaElement);
    553:   VecDestroy(&vecRho);
    554:   DMDestroy(&daVelAvg);
    555:   DMDestroy(&daP);
    556:   DMDestroy(&daEtaCorner);
    557:   DMDestroy(&daEtaElement);
    558:   DMDestroy(&daRho);
    559:   VecDestroy(&velAvg);
    560:   DMDestroy(&dmVelAvg);
    561:   return 0;
    562: }

    564: /*TEST

    566:    test:
    567:       suffix: direct_umfpack
    568:       requires: suitesparse !complex
    569:       nsize: 1
    570:       args: -stag_grid_x 12 -stag_grid_y 7 -pc_type lu -pc_factor_mat_solver_type umfpack

    572:    test:
    573:       suffix: direct_mumps
    574:       requires: mumps !complex
    575:       nsize: 9
    576:       args: -stag_grid_x 13 -stag_grid_y 8 -pc_type lu -pc_factor_mat_solver_type mumps

    578: TEST*/