Actual source code: ex4.c

petsc-3.11.3 2019-06-26
Report Typos and Errors
  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+1; 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:      DMStagVecGetArrayDOF() 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*/