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);if (ierr) return ierr;

     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:   KSPSetFromOptions(ksp);
    112:   KSPSolve(ksp,b,x);
    113:   {
    114:     KSPConvergedReason reason;
    115:     KSPGetConvergedReason(ksp,&reason);
    116:     if (reason < 0) SETERRQ(PETSC_COMM_WORLD,PETSC_ERR_CONV_FAILED,"Linear solve failed");
    117:   }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    384: /* Here, we demonstrate getting coordinates from a vector by using DMStagStencil.
    385: This would usually be done with direct array access, though. */
    386: static PetscErrorCode PopulateCoefficientData(Ctx ctx)
    387: {
    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: {
    474:   DM             dmVelAvg;
    475:   Vec            velAvg;
    476:   DM             daVelAvg,daP,daEtaElement,daEtaCorner,daRho;
    477:   Vec            vecVelAvg,vecP,vecEtaElement,vecEtaCorner,vecRho;

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

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

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

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

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

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

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

    565: /*TEST

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

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

    579: TEST*/