Actual source code: ex6.c

  1: static const char help[] = "Simple 2D or 3D finite difference seismic wave propagation, using\n"
  2:                            "a staggered grid represented with DMStag objects.\n"
  3:                            "-dim <2,3>  : dimension of problem\n"
  4:                            "-nsteps <n> : number of timesteps\n"
  5:                            "-dt <dt>    : length of timestep\n"
  6:                            "\n";

  8: /*
  9: Reference:

 11: @article{Virieux1986,
 12:   title     = {{P-SV} Wave Propagation in Heterogeneous Media: {V}elocity-Stress Finite-Difference Method},
 13:   author    = {Virieux, Jean},
 14:   journal   = {Geophysics},
 15:   volume    = {51},
 16:   number    = {4},
 17:   pages     = {889--901},
 18:   publisher = {Society of Exploration Geophysicists},
 19:   year      = {1986},
 20: }

 22: This example uses y in 2 dimensions, where the paper uses z.

 24: This example uses the dual grid of the one pictured in Fig. 1. of the paper, so
 25: that velocities are on face boundaries, shear stresses are defined on
 26: vertices(2D) or edges(3D), and normal stresses are defined on elements.

 28: There is a typo in the paragraph after (5) in the paper: Sigma, Xi, and Tau
 29: represent tau_xx, tau_xz, and tau_zz, respectively (the last two entries are
 30: transposed in the paper).

 32: This example treats the boundaries naively (by leaving ~zero velocity and
 33: stress there).
 34: */

 36: #include <petscdmstag.h>
 37: #include <petscts.h>

 39: /* A struct defining the parameters of the problem */
 40: typedef struct {
 41:   DM          dm_velocity, dm_stress;
 42:   DM          dm_buoyancy, dm_lame;
 43:   Vec         buoyancy, lame;  /* Global, but for efficiency one could store local vectors */
 44:   PetscInt    dim;             /* 2 or 3 */
 45:   PetscScalar rho, lambda, mu; /* constant material properties */
 46:   PetscReal   xmin, xmax, ymin, ymax, zmin, zmax;
 47:   PetscReal   dt;
 48:   PetscInt    timesteps;
 49:   PetscBool   dump_output;
 50: } Ctx;

 52: static PetscErrorCode CreateLame(Ctx *);
 53: static PetscErrorCode ForceStress(const Ctx *, Vec, PetscReal);
 54: static PetscErrorCode DumpVelocity(const Ctx *, Vec, PetscInt);
 55: static PetscErrorCode DumpStress(const Ctx *, Vec, PetscInt);
 56: static PetscErrorCode UpdateVelocity(const Ctx *, Vec, Vec, Vec);
 57: static PetscErrorCode UpdateStress(const Ctx *, Vec, Vec, Vec);

 59: int main(int argc, char *argv[])
 60: {
 61:   Ctx      ctx;
 62:   Vec      velocity, stress;
 63:   PetscInt timestep;

 65:   /* Initialize PETSc */
 67:   PetscInitialize(&argc, &argv, 0, help);

 69:   /* Populate application context */
 70:   ctx.dim         = 2;
 71:   ctx.rho         = 1.0;
 72:   ctx.lambda      = 1.0;
 73:   ctx.mu          = 1.0;
 74:   ctx.xmin        = 0.0;
 75:   ctx.xmax        = 1.0;
 76:   ctx.ymin        = 0.0;
 77:   ctx.ymax        = 1.0;
 78:   ctx.zmin        = 0.0;
 79:   ctx.zmax        = 1.0;
 80:   ctx.dt          = 0.001;
 81:   ctx.timesteps   = 100;
 82:   ctx.dump_output = PETSC_TRUE;

 84:   /* Update context from command line options */
 85:   PetscOptionsGetInt(NULL, NULL, "-dim", &ctx.dim, NULL);
 86:   PetscOptionsGetReal(NULL, NULL, "-dt", &ctx.dt, NULL);
 87:   PetscOptionsGetInt(NULL, NULL, "-nsteps", &ctx.timesteps, NULL);
 88:   PetscOptionsGetBool(NULL, NULL, "-dump_output", &ctx.dump_output, NULL);

 90:   /* Create a DMStag, with uniform coordinates, for the velocities */
 91:   {
 92:     PetscInt       dof0, dof1, dof2, dof3;
 93:     const PetscInt stencilWidth = 1;

 95:     switch (ctx.dim) {
 96:     case 2:
 97:       dof0 = 0;
 98:       dof1 = 1;
 99:       dof2 = 0; /* 1 dof per cell boundary */
100:       DMStagCreate2d(PETSC_COMM_WORLD, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, 100, 100, PETSC_DECIDE, PETSC_DECIDE, dof0, dof1, dof2, DMSTAG_STENCIL_BOX, stencilWidth, NULL, NULL, &ctx.dm_velocity);
101:       break;
102:     case 3:
103:       dof0 = 0;
104:       dof1 = 0;
105:       dof2 = 1;
106:       dof3 = 0; /* 1 dof per cell boundary */
107:       DMStagCreate3d(PETSC_COMM_WORLD, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, 30, 30, 30, PETSC_DECIDE, PETSC_DECIDE, PETSC_DECIDE, dof0, dof1, dof2, dof3, DMSTAG_STENCIL_BOX, stencilWidth, NULL, NULL, NULL, &ctx.dm_velocity);
108:       break;
109:     default:
110:       SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_SUP, "Not Implemented for dimension %" PetscInt_FMT, ctx.dim);
111:     }
112:   }
113:   DMSetFromOptions(ctx.dm_velocity); /* Options control velocity DM */
114:   DMSetUp(ctx.dm_velocity);
115:   DMStagSetUniformCoordinatesProduct(ctx.dm_velocity, ctx.xmin, ctx.xmax, ctx.ymin, ctx.ymax, ctx.zmin, ctx.zmax);

117:   /* Create a second, compatible DMStag for the stresses */
118:   switch (ctx.dim) {
119:   case 2:
120:     /* One shear stress component on element corners, two shear stress components on elements */
121:     DMStagCreateCompatibleDMStag(ctx.dm_velocity, 1, 0, 2, 0, &ctx.dm_stress);
122:     break;
123:   case 3:
124:     /* One shear stress component on element edges, three shear stress components on elements */
125:     DMStagCreateCompatibleDMStag(ctx.dm_velocity, 0, 1, 0, 3, &ctx.dm_stress);
126:     break;
127:   default:
128:     SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_SUP, "Not Implemented for dimension %" PetscInt_FMT, ctx.dim);
129:   }
130:   DMSetUp(ctx.dm_stress);
131:   DMStagSetUniformCoordinatesProduct(ctx.dm_stress, ctx.xmin, ctx.xmax, ctx.ymin, ctx.ymax, ctx.zmin, ctx.zmax);

133:   /* Create two additional DMStag objects for the buoyancy and Lame parameters */
134:   switch (ctx.dim) {
135:   case 2:
136:     /* buoyancy on element boundaries (edges) */
137:     DMStagCreateCompatibleDMStag(ctx.dm_velocity, 0, 1, 0, 0, &ctx.dm_buoyancy);
138:     break;
139:   case 3:
140:     /* buoyancy on element boundaries (faces) */
141:     DMStagCreateCompatibleDMStag(ctx.dm_velocity, 0, 0, 1, 0, &ctx.dm_buoyancy);
142:     break;
143:   default:
144:     SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_SUP, "Not Implemented for dimension %" PetscInt_FMT, ctx.dim);
145:   }
146:   DMSetUp(ctx.dm_buoyancy);

148:   switch (ctx.dim) {
149:   case 2:
150:     /* mu and lambda + 2*mu on element centers, mu on corners */
151:     DMStagCreateCompatibleDMStag(ctx.dm_velocity, 1, 0, 2, 0, &ctx.dm_lame);
152:     break;
153:   case 3:
154:     /* mu and lambda + 2*mu on element centers, mu on edges */
155:     DMStagCreateCompatibleDMStag(ctx.dm_velocity, 0, 1, 0, 2, &ctx.dm_lame);
156:     break;
157:   default:
158:     SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_SUP, "Not Implemented for dimension %" PetscInt_FMT, ctx.dim);
159:   }
160:   DMSetUp(ctx.dm_lame);

162:   /* Print out some info */
163:   {
164:     PetscInt    N[3];
165:     PetscScalar dx, Vp;

167:     DMStagGetGlobalSizes(ctx.dm_velocity, &N[0], &N[1], &N[2]);
168:     dx = (ctx.xmax - ctx.xmin) / N[0];
169:     Vp = PetscSqrtScalar((ctx.lambda + 2 * ctx.mu) / ctx.rho);
170:     if (ctx.dim == 2) {
171:       PetscPrintf(PETSC_COMM_WORLD, "Using a %" PetscInt_FMT " x %" PetscInt_FMT " mesh\n", N[0], N[1]);
172:     } else if (ctx.dim == 3) {
173:       PetscPrintf(PETSC_COMM_WORLD, "Using a %" PetscInt_FMT " x %" PetscInt_FMT " x %" PetscInt_FMT " mesh\n", N[0], N[1], N[2]);
174:     }
175:     PetscPrintf(PETSC_COMM_WORLD, "dx: %g\n", (double)PetscRealPart(dx));
176:     PetscPrintf(PETSC_COMM_WORLD, "dt: %g\n", (double)ctx.dt);
177:     PetscPrintf(PETSC_COMM_WORLD, "P-wave velocity: %g\n", (double)PetscRealPart(PetscSqrtScalar((ctx.lambda + 2 * ctx.mu) / ctx.rho)));
178:     PetscPrintf(PETSC_COMM_WORLD, "V_p dt / dx: %g\n", (double)PetscRealPart(Vp * ctx.dt / dx));
179:   }

181:   /* Populate the coefficient arrays */
182:   CreateLame(&ctx);
183:   DMCreateGlobalVector(ctx.dm_buoyancy, &ctx.buoyancy);
184:   VecSet(ctx.buoyancy, 1.0 / ctx.rho);

186:   /* Create vectors to store the system state */
187:   DMCreateGlobalVector(ctx.dm_velocity, &velocity);
188:   DMCreateGlobalVector(ctx.dm_stress, &stress);

190:   /* Initial State */
191:   VecSet(velocity, 0.0);
192:   VecSet(stress, 0.0);
193:   ForceStress(&ctx, stress, 0.0);
194:   if (ctx.dump_output) {
195:     DumpVelocity(&ctx, velocity, 0);
196:     DumpStress(&ctx, stress, 0);
197:   }

199:   /* Time Loop */
200:   for (timestep = 1; timestep <= ctx.timesteps; ++timestep) {
201:     const PetscReal t = timestep * ctx.dt;

203:     UpdateVelocity(&ctx, velocity, stress, ctx.buoyancy);
204:     UpdateStress(&ctx, velocity, stress, ctx.lame);
205:     ForceStress(&ctx, stress, t);
206:     PetscPrintf(PETSC_COMM_WORLD, "Timestep %" PetscInt_FMT ", t = %g\n", timestep, (double)t);
207:     if (ctx.dump_output) {
208:       DumpVelocity(&ctx, velocity, timestep);
209:       DumpStress(&ctx, stress, timestep);
210:     }
211:   }

213:   /* Clean up and finalize PETSc */
214:   VecDestroy(&velocity);
215:   VecDestroy(&stress);
216:   VecDestroy(&ctx.lame);
217:   VecDestroy(&ctx.buoyancy);
218:   DMDestroy(&ctx.dm_velocity);
219:   DMDestroy(&ctx.dm_stress);
220:   DMDestroy(&ctx.dm_buoyancy);
221:   DMDestroy(&ctx.dm_lame);
222:   PetscFinalize();
223:   return 0;
224: }

226: static PetscErrorCode CreateLame(Ctx *ctx)
227: {
228:   PetscInt N[3], ex, ey, ez, startx, starty, startz, nx, ny, nz, extrax, extray, extraz;

231:   DMCreateGlobalVector(ctx->dm_lame, &ctx->lame);
232:   DMStagGetGlobalSizes(ctx->dm_lame, &N[0], &N[1], &N[2]);
233:   DMStagGetCorners(ctx->dm_buoyancy, &startx, &starty, &startz, &nx, &ny, &nz, &extrax, &extray, &extraz);
234:   if (ctx->dim == 2) {
235:     /* Element values */
236:     for (ey = starty; ey < starty + ny; ++ey) {
237:       for (ex = startx; ex < startx + nx; ++ex) {
238:         DMStagStencil pos;

240:         pos.i   = ex;
241:         pos.j   = ey;
242:         pos.c   = 0;
243:         pos.loc = DMSTAG_ELEMENT;
244:         DMStagVecSetValuesStencil(ctx->dm_lame, ctx->lame, 1, &pos, &ctx->lambda, INSERT_VALUES);
245:         pos.i   = ex;
246:         pos.j   = ey;
247:         pos.c   = 1;
248:         pos.loc = DMSTAG_ELEMENT;
249:         DMStagVecSetValuesStencil(ctx->dm_lame, ctx->lame, 1, &pos, &ctx->mu, INSERT_VALUES);
250:       }
251:     }
252:     /* Vertex Values */
253:     for (ey = starty; ey < starty + ny + extray; ++ey) {
254:       for (ex = startx; ex < startx + nx + extrax; ++ex) {
255:         DMStagStencil pos;

257:         pos.i   = ex;
258:         pos.j   = ey;
259:         pos.c   = 0;
260:         pos.loc = DMSTAG_DOWN_LEFT;
261:         DMStagVecSetValuesStencil(ctx->dm_lame, ctx->lame, 1, &pos, &ctx->mu, INSERT_VALUES);
262:       }
263:     }
264:   } else if (ctx->dim == 3) {
265:     /* Element values */
266:     for (ez = startz; ez < startz + nz; ++ez) {
267:       for (ey = starty; ey < starty + ny; ++ey) {
268:         for (ex = startx; ex < startx + nx; ++ex) {
269:           DMStagStencil pos;

271:           pos.i   = ex;
272:           pos.j   = ey;
273:           pos.k   = ez;
274:           pos.c   = 0;
275:           pos.loc = DMSTAG_ELEMENT;
276:           DMStagVecSetValuesStencil(ctx->dm_lame, ctx->lame, 1, &pos, &ctx->lambda, INSERT_VALUES);
277:           pos.i   = ex;
278:           pos.j   = ey;
279:           pos.k   = ez;
280:           pos.c   = 1;
281:           pos.loc = DMSTAG_ELEMENT;
282:           DMStagVecSetValuesStencil(ctx->dm_lame, ctx->lame, 1, &pos, &ctx->mu, INSERT_VALUES);
283:         }
284:       }
285:     }
286:     /* Edge Values */
287:     for (ez = startz; ez < startz + nz + extraz; ++ez) {
288:       for (ey = starty; ey < starty + ny + extray; ++ey) {
289:         for (ex = startx; ex < startx + nx + extrax; ++ex) {
290:           DMStagStencil pos;

292:           if (ex < N[0]) {
293:             pos.i   = ex;
294:             pos.j   = ey;
295:             pos.k   = ez;
296:             pos.c   = 0;
297:             pos.loc = DMSTAG_BACK_DOWN;
298:             DMStagVecSetValuesStencil(ctx->dm_lame, ctx->lame, 1, &pos, &ctx->mu, INSERT_VALUES);
299:           }
300:           if (ey < N[1]) {
301:             pos.i   = ex;
302:             pos.j   = ey;
303:             pos.k   = ez;
304:             pos.c   = 0;
305:             pos.loc = DMSTAG_BACK_LEFT;
306:             DMStagVecSetValuesStencil(ctx->dm_lame, ctx->lame, 1, &pos, &ctx->mu, INSERT_VALUES);
307:           }
308:           if (ez < N[2]) {
309:             pos.i   = ex;
310:             pos.j   = ey;
311:             pos.k   = ez;
312:             pos.c   = 0;
313:             pos.loc = DMSTAG_DOWN_LEFT;
314:             DMStagVecSetValuesStencil(ctx->dm_lame, ctx->lame, 1, &pos, &ctx->mu, INSERT_VALUES);
315:           }
316:         }
317:       }
318:     }
319:   } else SETERRQ(PetscObjectComm((PetscObject)ctx->dm_velocity), PETSC_ERR_SUP, "Unsupported dim %" PetscInt_FMT, ctx->dim);
320:   VecAssemblyBegin(ctx->lame);
321:   VecAssemblyEnd(ctx->lame);
322:   return 0;
323: }

325: static PetscErrorCode ForceStress(const Ctx *ctx, Vec stress, PetscReal t)
326: {
327:   PetscInt          start[3], n[3], N[3];
328:   DMStagStencil     pos;
329:   PetscBool         this_rank;
330:   const PetscScalar val = PetscExpReal(-100.0 * t);

333:   DMStagGetGlobalSizes(ctx->dm_stress, &N[0], &N[1], &N[2]);
334:   DMStagGetCorners(ctx->dm_stress, &start[0], &start[1], &start[2], &n[0], &n[1], &n[2], NULL, NULL, NULL);

336:   /* Normal stresses at a single point */
337:   this_rank = (PetscBool)(start[0] <= N[0] / 2 && N[0] / 2 <= start[0] + n[0]);
338:   this_rank = (PetscBool)(this_rank && start[1] <= N[1] / 2 && N[1] / 2 <= start[1] + n[1]);
339:   if (ctx->dim == 3) this_rank = (PetscBool)(this_rank && start[2] <= N[2] / 2 && N[2] / 2 <= start[2] + n[2]);
340:   if (this_rank) {
341:     /* Note integer division to pick element near the center */
342:     pos.i   = N[0] / 2;
343:     pos.j   = N[1] / 2;
344:     pos.k   = N[2] / 2;
345:     pos.c   = 0;
346:     pos.loc = DMSTAG_ELEMENT;
347:     DMStagVecSetValuesStencil(ctx->dm_stress, stress, 1, &pos, &val, INSERT_VALUES);
348:     pos.i   = N[0] / 2;
349:     pos.j   = N[1] / 2;
350:     pos.k   = N[2] / 2;
351:     pos.c   = 1;
352:     pos.loc = DMSTAG_ELEMENT;
353:     DMStagVecSetValuesStencil(ctx->dm_stress, stress, 1, &pos, &val, INSERT_VALUES);
354:     if (ctx->dim == 3) {
355:       pos.i   = N[0] / 2;
356:       pos.j   = N[1] / 2;
357:       pos.k   = N[2] / 2;
358:       pos.c   = 2;
359:       pos.loc = DMSTAG_ELEMENT;
360:       DMStagVecSetValuesStencil(ctx->dm_stress, stress, 1, &pos, &val, INSERT_VALUES);
361:     }
362:   }

364:   VecAssemblyBegin(stress);
365:   VecAssemblyEnd(stress);
366:   return 0;
367: }

369: static PetscErrorCode UpdateVelocity_2d(const Ctx *ctx, Vec velocity, Vec stress, Vec buoyancy)
370: {
371:   Vec                  velocity_local, stress_local, buoyancy_local;
372:   PetscInt             ex, ey, startx, starty, nx, ny;
373:   PetscInt             slot_coord_next, slot_coord_element, slot_coord_prev;
374:   PetscInt             slot_vx_left, slot_vy_down, slot_buoyancy_down, slot_buoyancy_left;
375:   PetscInt             slot_txx, slot_tyy, slot_txy_downleft, slot_txy_downright, slot_txy_upleft;
376:   const PetscScalar  **arr_coord_x, **arr_coord_y;
377:   const PetscScalar ***arr_stress, ***arr_buoyancy;
378:   PetscScalar       ***arr_velocity;


382:   /* Prepare direct access to buoyancy data */
383:   DMStagGetLocationSlot(ctx->dm_buoyancy, DMSTAG_LEFT, 0, &slot_buoyancy_left);
384:   DMStagGetLocationSlot(ctx->dm_buoyancy, DMSTAG_DOWN, 0, &slot_buoyancy_down);
385:   DMGetLocalVector(ctx->dm_buoyancy, &buoyancy_local);
386:   DMGlobalToLocal(ctx->dm_buoyancy, buoyancy, INSERT_VALUES, buoyancy_local);
387:   DMStagVecGetArrayRead(ctx->dm_buoyancy, buoyancy_local, (void *)&arr_buoyancy);

389:   /* Prepare read-only access to stress data */
390:   DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_ELEMENT, 0, &slot_txx);
391:   DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_ELEMENT, 1, &slot_tyy);
392:   DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_UP_LEFT, 0, &slot_txy_upleft);
393:   DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_DOWN_LEFT, 0, &slot_txy_downleft);
394:   DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_DOWN_RIGHT, 0, &slot_txy_downright);
395:   DMGetLocalVector(ctx->dm_stress, &stress_local);
396:   DMGlobalToLocal(ctx->dm_stress, stress, INSERT_VALUES, stress_local);
397:   DMStagVecGetArrayRead(ctx->dm_stress, stress_local, (void *)&arr_stress);

399:   /* Prepare read-write access to velocity data */
400:   DMStagGetLocationSlot(ctx->dm_velocity, DMSTAG_LEFT, 0, &slot_vx_left);
401:   DMStagGetLocationSlot(ctx->dm_velocity, DMSTAG_DOWN, 0, &slot_vy_down);
402:   DMGetLocalVector(ctx->dm_velocity, &velocity_local);
403:   DMGlobalToLocal(ctx->dm_velocity, velocity, INSERT_VALUES, velocity_local);
404:   DMStagVecGetArray(ctx->dm_velocity, velocity_local, &arr_velocity);

406:   /* Prepare read-only access to coordinate data */
407:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_LEFT, &slot_coord_prev);
408:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_RIGHT, &slot_coord_next);
409:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_ELEMENT, &slot_coord_element);
410:   DMStagGetProductCoordinateArrays(ctx->dm_velocity, (void *)&arr_coord_x, (void *)&arr_coord_y, NULL);

412:   /* Iterate over interior of the domain, updating the velocities */
413:   DMStagGetCorners(ctx->dm_velocity, &startx, &starty, NULL, &nx, &ny, NULL, NULL, NULL, NULL);
414:   for (ey = starty; ey < starty + ny; ++ey) {
415:     for (ex = startx; ex < startx + nx; ++ex) {
416:       /* Update y-velocity */
417:       if (ey > 0) {
418:         const PetscScalar dx = arr_coord_x[ex][slot_coord_next] - arr_coord_x[ex][slot_coord_prev];
419:         const PetscScalar dy = arr_coord_y[ey][slot_coord_element] - arr_coord_y[ey - 1][slot_coord_element];
420:         const PetscScalar B  = arr_buoyancy[ey][ex][slot_buoyancy_down];

422:         arr_velocity[ey][ex][slot_vy_down] += B * ctx->dt * ((arr_stress[ey][ex][slot_txy_downright] - arr_stress[ey][ex][slot_txy_downleft]) / dx + (arr_stress[ey][ex][slot_tyy] - arr_stress[ey - 1][ex][slot_tyy]) / dy);
423:       }

425:       /* Update x-velocity */
426:       if (ex > 0) {
427:         const PetscScalar dx = arr_coord_x[ex][slot_coord_element] - arr_coord_x[ex - 1][slot_coord_element];
428:         const PetscScalar dy = arr_coord_y[ey][slot_coord_next] - arr_coord_y[ey][slot_coord_prev];
429:         const PetscScalar B  = arr_buoyancy[ey][ex][slot_buoyancy_left];

431:         arr_velocity[ey][ex][slot_vx_left] += B * ctx->dt * ((arr_stress[ey][ex][slot_txx] - arr_stress[ey][ex - 1][slot_txx]) / dx + (arr_stress[ey][ex][slot_txy_upleft] - arr_stress[ey][ex][slot_txy_downleft]) / dy);
432:       }
433:     }
434:   }

436:   /* Restore all access */
437:   DMStagVecRestoreArrayRead(ctx->dm_buoyancy, buoyancy_local, (void *)&arr_buoyancy);
438:   DMRestoreLocalVector(ctx->dm_buoyancy, &buoyancy_local);
439:   DMStagVecRestoreArrayRead(ctx->dm_stress, stress_local, (void *)&arr_stress);
440:   DMRestoreLocalVector(ctx->dm_stress, &stress_local);
441:   DMStagVecRestoreArray(ctx->dm_velocity, velocity_local, &arr_velocity);
442:   DMLocalToGlobal(ctx->dm_velocity, velocity_local, INSERT_VALUES, velocity);
443:   DMRestoreLocalVector(ctx->dm_velocity, &velocity_local);
444:   DMStagRestoreProductCoordinateArrays(ctx->dm_velocity, (void *)&arr_coord_x, (void *)&arr_coord_y, NULL);
445:   return 0;
446: }

448: static PetscErrorCode UpdateVelocity_3d(const Ctx *ctx, Vec velocity, Vec stress, Vec buoyancy)
449: {
450:   Vec                   velocity_local, stress_local, buoyancy_local;
451:   PetscInt              ex, ey, ez, startx, starty, startz, nx, ny, nz;
452:   PetscInt              slot_coord_next, slot_coord_element, slot_coord_prev;
453:   PetscInt              slot_vx_left, slot_vy_down, slot_vz_back, slot_buoyancy_down, slot_buoyancy_left, slot_buoyancy_back;
454:   PetscInt              slot_txx, slot_tyy, slot_tzz;
455:   PetscInt              slot_txy_downleft, slot_txy_downright, slot_txy_upleft;
456:   PetscInt              slot_txz_backleft, slot_txz_backright, slot_txz_frontleft;
457:   PetscInt              slot_tyz_backdown, slot_tyz_frontdown, slot_tyz_backup;
458:   const PetscScalar   **arr_coord_x, **arr_coord_y, **arr_coord_z;
459:   const PetscScalar ****arr_stress, ****arr_buoyancy;
460:   PetscScalar       ****arr_velocity;


464:   /* Prepare direct access to buoyancy data */
465:   DMStagGetLocationSlot(ctx->dm_buoyancy, DMSTAG_LEFT, 0, &slot_buoyancy_left);
466:   DMStagGetLocationSlot(ctx->dm_buoyancy, DMSTAG_DOWN, 0, &slot_buoyancy_down);
467:   DMStagGetLocationSlot(ctx->dm_buoyancy, DMSTAG_BACK, 0, &slot_buoyancy_back);
468:   DMGetLocalVector(ctx->dm_buoyancy, &buoyancy_local);
469:   DMGlobalToLocal(ctx->dm_buoyancy, buoyancy, INSERT_VALUES, buoyancy_local);
470:   DMStagVecGetArrayRead(ctx->dm_buoyancy, buoyancy_local, (void *)&arr_buoyancy);

472:   /* Prepare read-only access to stress data */
473:   DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_ELEMENT, 0, &slot_txx);
474:   DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_ELEMENT, 1, &slot_tyy);
475:   DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_ELEMENT, 2, &slot_tzz);
476:   DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_UP_LEFT, 0, &slot_txy_upleft);
477:   DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_DOWN_LEFT, 0, &slot_txy_downleft);
478:   DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_DOWN_RIGHT, 0, &slot_txy_downright);
479:   DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_BACK_LEFT, 0, &slot_txz_backleft);
480:   DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_BACK_RIGHT, 0, &slot_txz_backright);
481:   DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_FRONT_LEFT, 0, &slot_txz_frontleft);
482:   DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_BACK_DOWN, 0, &slot_tyz_backdown);
483:   DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_BACK_UP, 0, &slot_tyz_backup);
484:   DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_FRONT_DOWN, 0, &slot_tyz_frontdown);
485:   DMGetLocalVector(ctx->dm_stress, &stress_local);
486:   DMGlobalToLocal(ctx->dm_stress, stress, INSERT_VALUES, stress_local);
487:   DMStagVecGetArrayRead(ctx->dm_stress, stress_local, (void *)&arr_stress);

489:   /* Prepare read-write access to velocity data */
490:   DMStagGetLocationSlot(ctx->dm_velocity, DMSTAG_LEFT, 0, &slot_vx_left);
491:   DMStagGetLocationSlot(ctx->dm_velocity, DMSTAG_DOWN, 0, &slot_vy_down);
492:   DMStagGetLocationSlot(ctx->dm_velocity, DMSTAG_BACK, 0, &slot_vz_back);
493:   DMGetLocalVector(ctx->dm_velocity, &velocity_local);
494:   DMGlobalToLocal(ctx->dm_velocity, velocity, INSERT_VALUES, velocity_local);
495:   DMStagVecGetArray(ctx->dm_velocity, velocity_local, &arr_velocity);

497:   /* Prepare read-only access to coordinate data */
498:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_LEFT, &slot_coord_prev);
499:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_RIGHT, &slot_coord_next);
500:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_ELEMENT, &slot_coord_element);
501:   DMStagGetProductCoordinateArrays(ctx->dm_velocity, (void *)&arr_coord_x, (void *)&arr_coord_y, (void *)&arr_coord_z);

503:   /* Iterate over interior of the domain, updating the velocities */
504:   DMStagGetCorners(ctx->dm_velocity, &startx, &starty, &startz, &nx, &ny, &nz, NULL, NULL, NULL);
505:   for (ez = startz; ez < startz + nz; ++ez) {
506:     for (ey = starty; ey < starty + ny; ++ey) {
507:       for (ex = startx; ex < startx + nx; ++ex) {
508:         /* Update y-velocity */
509:         if (ey > 0) {
510:           const PetscScalar dx = arr_coord_x[ex][slot_coord_next] - arr_coord_x[ex][slot_coord_prev];
511:           const PetscScalar dy = arr_coord_y[ey][slot_coord_element] - arr_coord_y[ey - 1][slot_coord_element];
512:           const PetscScalar dz = arr_coord_z[ez][slot_coord_next] - arr_coord_z[ez][slot_coord_prev];
513:           const PetscScalar B  = arr_buoyancy[ez][ey][ex][slot_buoyancy_down];

515:           arr_velocity[ez][ey][ex][slot_vy_down] += B * ctx->dt * ((arr_stress[ez][ey][ex][slot_txy_downright] - arr_stress[ez][ey][ex][slot_txy_downleft]) / dx + (arr_stress[ez][ey][ex][slot_tyy] - arr_stress[ez][ey - 1][ex][slot_tyy]) / dy + (arr_stress[ez][ey][ex][slot_tyz_frontdown] - arr_stress[ez][ey][ex][slot_tyz_backdown]) / dz);
516:         }

518:         /* Update x-velocity */
519:         if (ex > 0) {
520:           const PetscScalar dx = arr_coord_x[ex][slot_coord_element] - arr_coord_x[ex - 1][slot_coord_element];
521:           const PetscScalar dy = arr_coord_y[ey][slot_coord_next] - arr_coord_y[ey][slot_coord_prev];
522:           const PetscScalar dz = arr_coord_z[ez][slot_coord_next] - arr_coord_z[ez][slot_coord_prev];
523:           const PetscScalar B  = arr_buoyancy[ez][ey][ex][slot_buoyancy_left];

525:           arr_velocity[ez][ey][ex][slot_vx_left] += B * ctx->dt * ((arr_stress[ez][ey][ex][slot_txx] - arr_stress[ez][ey][ex - 1][slot_txx]) / dx + (arr_stress[ez][ey][ex][slot_txy_upleft] - arr_stress[ez][ey][ex][slot_txy_downleft]) / dy + (arr_stress[ez][ey][ex][slot_txz_frontleft] - arr_stress[ez][ey][ex][slot_txz_backleft]) / dz);
526:         }

528:         /* Update z-velocity */
529:         if (ez > 0) {
530:           const PetscScalar dx = arr_coord_x[ex][slot_coord_next] - arr_coord_x[ex][slot_coord_prev];
531:           const PetscScalar dy = arr_coord_y[ey][slot_coord_next] - arr_coord_y[ey][slot_coord_prev];
532:           const PetscScalar dz = arr_coord_z[ez][slot_coord_element] - arr_coord_z[ez - 1][slot_coord_element];
533:           const PetscScalar B  = arr_buoyancy[ez][ey][ex][slot_buoyancy_back];

535:           arr_velocity[ez][ey][ex][slot_vz_back] += B * ctx->dt * ((arr_stress[ez][ey][ex][slot_txz_backright] - arr_stress[ez][ey][ex][slot_txz_backleft]) / dx + (arr_stress[ez][ey][ex][slot_tyz_backup] - arr_stress[ez][ey][ex][slot_tyz_backdown]) / dy + (arr_stress[ez][ey][ex][slot_tzz] - arr_stress[ez - 1][ey][ex][slot_tzz]) / dz);
536:         }
537:       }
538:     }
539:   }

541:   /* Restore all access */
542:   DMStagVecRestoreArrayRead(ctx->dm_buoyancy, buoyancy_local, (void *)&arr_buoyancy);
543:   DMRestoreLocalVector(ctx->dm_buoyancy, &buoyancy_local);
544:   DMStagVecRestoreArrayRead(ctx->dm_stress, stress_local, (void *)&arr_stress);
545:   DMRestoreLocalVector(ctx->dm_stress, &stress_local);
546:   DMStagVecRestoreArray(ctx->dm_velocity, velocity_local, &arr_velocity);
547:   DMLocalToGlobal(ctx->dm_velocity, velocity_local, INSERT_VALUES, velocity);
548:   DMRestoreLocalVector(ctx->dm_velocity, &velocity_local);
549:   DMStagRestoreProductCoordinateArrays(ctx->dm_velocity, (void *)&arr_coord_x, (void *)&arr_coord_y, (void *)&arr_coord_z);
550:   return 0;
551: }

553: static PetscErrorCode UpdateVelocity(const Ctx *ctx, Vec velocity, Vec stress, Vec buoyancy)
554: {
556:   if (ctx->dim == 2) {
557:     UpdateVelocity_2d(ctx, velocity, stress, buoyancy);
558:   } else if (ctx->dim == 3) {
559:     UpdateVelocity_3d(ctx, velocity, stress, buoyancy);
560:   } else SETERRQ(PetscObjectComm((PetscObject)ctx->dm_velocity), PETSC_ERR_SUP, "Unsupported dim %" PetscInt_FMT, ctx->dim);
561:   return 0;
562: }

564: static PetscErrorCode UpdateStress_2d(const Ctx *ctx, Vec velocity, Vec stress, Vec lame)
565: {
566:   Vec                  velocity_local, stress_local, lame_local;
567:   PetscInt             ex, ey, startx, starty, nx, ny;
568:   PetscInt             slot_coord_next, slot_coord_element, slot_coord_prev;
569:   PetscInt             slot_vx_left, slot_vy_down, slot_vx_right, slot_vy_up;
570:   PetscInt             slot_mu_element, slot_lambda_element, slot_mu_downleft;
571:   PetscInt             slot_txx, slot_tyy, slot_txy_downleft;
572:   const PetscScalar  **arr_coord_x, **arr_coord_y;
573:   const PetscScalar ***arr_velocity, ***arr_lame;
574:   PetscScalar       ***arr_stress;


578:   /* Prepare read-write access to stress data */
579:   DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_ELEMENT, 0, &slot_txx);
580:   DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_ELEMENT, 1, &slot_tyy);
581:   DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_DOWN_LEFT, 0, &slot_txy_downleft);
582:   DMGetLocalVector(ctx->dm_stress, &stress_local);
583:   DMGlobalToLocal(ctx->dm_stress, stress, INSERT_VALUES, stress_local);
584:   DMStagVecGetArray(ctx->dm_stress, stress_local, &arr_stress);

586:   /* Prepare read-only access to velocity data */
587:   DMStagGetLocationSlot(ctx->dm_velocity, DMSTAG_LEFT, 0, &slot_vx_left);
588:   DMStagGetLocationSlot(ctx->dm_velocity, DMSTAG_RIGHT, 0, &slot_vx_right);
589:   DMStagGetLocationSlot(ctx->dm_velocity, DMSTAG_DOWN, 0, &slot_vy_down);
590:   DMStagGetLocationSlot(ctx->dm_velocity, DMSTAG_UP, 0, &slot_vy_up);
591:   DMGetLocalVector(ctx->dm_velocity, &velocity_local);
592:   DMGlobalToLocal(ctx->dm_velocity, velocity, INSERT_VALUES, velocity_local);
593:   DMStagVecGetArrayRead(ctx->dm_velocity, velocity_local, (void *)&arr_velocity);

595:   /* Prepare read-only access to Lame' data */
596:   DMStagGetLocationSlot(ctx->dm_lame, DMSTAG_ELEMENT, 0, &slot_lambda_element);
597:   DMStagGetLocationSlot(ctx->dm_lame, DMSTAG_ELEMENT, 1, &slot_mu_element);
598:   DMStagGetLocationSlot(ctx->dm_lame, DMSTAG_DOWN_LEFT, 0, &slot_mu_downleft);
599:   DMGetLocalVector(ctx->dm_lame, &lame_local);
600:   DMGlobalToLocal(ctx->dm_lame, lame, INSERT_VALUES, lame_local);
601:   DMStagVecGetArrayRead(ctx->dm_lame, lame_local, (void *)&arr_lame);

603:   /* Prepare read-only access to coordinate data */
604:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_LEFT, &slot_coord_prev);
605:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_RIGHT, &slot_coord_next);
606:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_ELEMENT, &slot_coord_element);
607:   DMStagGetProductCoordinateArrays(ctx->dm_velocity, (void *)&arr_coord_x, (void *)&arr_coord_y, NULL);

609:   /* Iterate over the interior of the domain, updating the stresses */
610:   DMStagGetCorners(ctx->dm_velocity, &startx, &starty, NULL, &nx, &ny, NULL, NULL, NULL, NULL);
611:   for (ey = starty; ey < starty + ny; ++ey) {
612:     for (ex = startx; ex < startx + nx; ++ex) {
613:       /* Update tau_xx and tau_yy*/
614:       {
615:         const PetscScalar dx   = arr_coord_x[ex][slot_coord_next] - arr_coord_x[ex][slot_coord_prev];
616:         const PetscScalar dy   = arr_coord_y[ey][slot_coord_next] - arr_coord_y[ey][slot_coord_prev];
617:         const PetscScalar L    = arr_lame[ey][ex][slot_lambda_element];
618:         const PetscScalar Lp2M = arr_lame[ey][ex][slot_lambda_element] + 2 * arr_lame[ey][ex][slot_mu_element];

620:         arr_stress[ey][ex][slot_txx] += Lp2M * ctx->dt * (arr_velocity[ey][ex][slot_vx_right] - arr_velocity[ey][ex][slot_vx_left]) / dx + L * ctx->dt * (arr_velocity[ey][ex][slot_vy_up] - arr_velocity[ey][ex][slot_vy_down]) / dy;

622:         arr_stress[ey][ex][slot_tyy] += Lp2M * ctx->dt * (arr_velocity[ey][ex][slot_vy_up] - arr_velocity[ey][ex][slot_vy_down]) / dy + L * ctx->dt * (arr_velocity[ey][ex][slot_vx_right] - arr_velocity[ey][ex][slot_vx_left]) / dx;
623:       }

625:       /* Update tau_xy */
626:       if (ex > 0 && ey > 0) {
627:         const PetscScalar dx = arr_coord_x[ex][slot_coord_element] - arr_coord_x[ex - 1][slot_coord_element];
628:         const PetscScalar dy = arr_coord_y[ey][slot_coord_element] - arr_coord_y[ey - 1][slot_coord_element];
629:         const PetscScalar M  = arr_lame[ey][ex][slot_mu_downleft];

631:         arr_stress[ey][ex][slot_txy_downleft] += M * ctx->dt * ((arr_velocity[ey][ex][slot_vx_left] - arr_velocity[ey - 1][ex][slot_vx_left]) / dy + (arr_velocity[ey][ex][slot_vy_down] - arr_velocity[ey][ex - 1][slot_vy_down]) / dx);
632:       }
633:     }
634:   }

636:   /* Restore all access */
637:   DMStagVecRestoreArray(ctx->dm_stress, stress_local, &arr_stress);
638:   DMLocalToGlobal(ctx->dm_stress, stress_local, INSERT_VALUES, stress);
639:   DMRestoreLocalVector(ctx->dm_stress, &stress_local);
640:   DMStagVecRestoreArrayRead(ctx->dm_velocity, velocity_local, (void *)&arr_velocity);
641:   DMRestoreLocalVector(ctx->dm_velocity, &velocity_local);
642:   DMStagVecRestoreArrayRead(ctx->dm_lame, lame_local, (void *)&arr_lame);
643:   DMRestoreLocalVector(ctx->dm_lame, &lame_local);
644:   DMStagRestoreProductCoordinateArrays(ctx->dm_velocity, (void *)&arr_coord_x, (void *)&arr_coord_y, NULL);
645:   return 0;
646: }

648: static PetscErrorCode UpdateStress_3d(const Ctx *ctx, Vec velocity, Vec stress, Vec lame)
649: {
650:   Vec                   velocity_local, stress_local, lame_local;
651:   PetscInt              ex, ey, ez, startx, starty, startz, nx, ny, nz;
652:   PetscInt              slot_coord_next, slot_coord_element, slot_coord_prev;
653:   PetscInt              slot_vx_left, slot_vy_down, slot_vx_right, slot_vy_up, slot_vz_back, slot_vz_front;
654:   PetscInt              slot_mu_element, slot_lambda_element, slot_mu_downleft, slot_mu_backdown, slot_mu_backleft;
655:   PetscInt              slot_txx, slot_tyy, slot_tzz, slot_txy_downleft, slot_txz_backleft, slot_tyz_backdown;
656:   const PetscScalar   **arr_coord_x, **arr_coord_y, **arr_coord_z;
657:   const PetscScalar ****arr_velocity, ****arr_lame;
658:   PetscScalar       ****arr_stress;


662:   /* Prepare read-write access to stress data */
663:   DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_ELEMENT, 0, &slot_txx);
664:   DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_ELEMENT, 1, &slot_tyy);
665:   DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_ELEMENT, 2, &slot_tzz);
666:   DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_DOWN_LEFT, 0, &slot_txy_downleft);
667:   DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_BACK_LEFT, 0, &slot_txz_backleft);
668:   DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_BACK_DOWN, 0, &slot_tyz_backdown);
669:   DMGetLocalVector(ctx->dm_stress, &stress_local);
670:   DMGlobalToLocal(ctx->dm_stress, stress, INSERT_VALUES, stress_local);
671:   DMStagVecGetArray(ctx->dm_stress, stress_local, &arr_stress);

673:   /* Prepare read-only access to velocity data */
674:   DMStagGetLocationSlot(ctx->dm_velocity, DMSTAG_LEFT, 0, &slot_vx_left);
675:   DMStagGetLocationSlot(ctx->dm_velocity, DMSTAG_RIGHT, 0, &slot_vx_right);
676:   DMStagGetLocationSlot(ctx->dm_velocity, DMSTAG_DOWN, 0, &slot_vy_down);
677:   DMStagGetLocationSlot(ctx->dm_velocity, DMSTAG_UP, 0, &slot_vy_up);
678:   DMStagGetLocationSlot(ctx->dm_velocity, DMSTAG_BACK, 0, &slot_vz_back);
679:   DMStagGetLocationSlot(ctx->dm_velocity, DMSTAG_FRONT, 0, &slot_vz_front);
680:   DMGetLocalVector(ctx->dm_velocity, &velocity_local);
681:   DMGlobalToLocal(ctx->dm_velocity, velocity, INSERT_VALUES, velocity_local);
682:   DMStagVecGetArrayRead(ctx->dm_velocity, velocity_local, (void *)&arr_velocity);

684:   /* Prepare read-only access to Lame' data */
685:   DMStagGetLocationSlot(ctx->dm_lame, DMSTAG_ELEMENT, 0, &slot_lambda_element);
686:   DMStagGetLocationSlot(ctx->dm_lame, DMSTAG_ELEMENT, 1, &slot_mu_element);
687:   DMStagGetLocationSlot(ctx->dm_lame, DMSTAG_DOWN_LEFT, 0, &slot_mu_downleft);
688:   DMStagGetLocationSlot(ctx->dm_lame, DMSTAG_BACK_LEFT, 0, &slot_mu_backleft);
689:   DMStagGetLocationSlot(ctx->dm_lame, DMSTAG_BACK_DOWN, 0, &slot_mu_backdown);
690:   DMGetLocalVector(ctx->dm_lame, &lame_local);
691:   DMGlobalToLocal(ctx->dm_lame, lame, INSERT_VALUES, lame_local);
692:   DMStagVecGetArrayRead(ctx->dm_lame, lame_local, (void *)&arr_lame);

694:   /* Prepare read-only access to coordinate data */
695:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_LEFT, &slot_coord_prev);
696:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_RIGHT, &slot_coord_next);
697:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_ELEMENT, &slot_coord_element);
698:   DMStagGetProductCoordinateArrays(ctx->dm_velocity, (void *)&arr_coord_x, (void *)&arr_coord_y, (void *)&arr_coord_z);

700:   /* Iterate over the interior of the domain, updating the stresses */
701:   DMStagGetCorners(ctx->dm_velocity, &startx, &starty, &startz, &nx, &ny, &nz, NULL, NULL, NULL);
702:   for (ez = startz; ez < startz + nz; ++ez) {
703:     for (ey = starty; ey < starty + ny; ++ey) {
704:       for (ex = startx; ex < startx + nx; ++ex) {
705:         /* Update tau_xx, tau_yy, and tau_zz*/
706:         {
707:           const PetscScalar dx   = arr_coord_x[ex][slot_coord_next] - arr_coord_x[ex][slot_coord_prev];
708:           const PetscScalar dy   = arr_coord_y[ey][slot_coord_next] - arr_coord_y[ey][slot_coord_prev];
709:           const PetscScalar dz   = arr_coord_z[ez][slot_coord_next] - arr_coord_z[ez][slot_coord_prev];
710:           const PetscScalar L    = arr_lame[ez][ey][ex][slot_lambda_element];
711:           const PetscScalar Lp2M = arr_lame[ez][ey][ex][slot_lambda_element] + 2 * arr_lame[ez][ey][ex][slot_mu_element];

713:           arr_stress[ez][ey][ex][slot_txx] += Lp2M * ctx->dt * (arr_velocity[ez][ey][ex][slot_vx_right] - arr_velocity[ez][ey][ex][slot_vx_left]) / dx + L * ctx->dt * (arr_velocity[ez][ey][ex][slot_vy_up] - arr_velocity[ez][ey][ex][slot_vy_down]) / dy +
714:                                               L * ctx->dt * (arr_velocity[ez][ey][ex][slot_vz_front] - arr_velocity[ez][ey][ex][slot_vz_back]) / dz;

716:           arr_stress[ez][ey][ex][slot_tyy] += Lp2M * ctx->dt * (arr_velocity[ez][ey][ex][slot_vy_up] - arr_velocity[ez][ey][ex][slot_vy_down]) / dy + L * ctx->dt * (arr_velocity[ez][ey][ex][slot_vz_front] - arr_velocity[ez][ey][ex][slot_vz_back]) / dz +
717:                                               L * ctx->dt * (arr_velocity[ez][ey][ex][slot_vx_right] - arr_velocity[ez][ey][ex][slot_vx_left]) / dx;

719:           arr_stress[ez][ey][ex][slot_tzz] += Lp2M * ctx->dt * (arr_velocity[ez][ey][ex][slot_vz_front] - arr_velocity[ez][ey][ex][slot_vz_back]) / dz + L * ctx->dt * (arr_velocity[ez][ey][ex][slot_vx_right] - arr_velocity[ez][ey][ex][slot_vx_left]) / dx +
720:                                               L * ctx->dt * (arr_velocity[ez][ey][ex][slot_vy_up] - arr_velocity[ez][ey][ex][slot_vy_down]) / dy;
721:         }

723:         /* Update tau_xy, tau_xz, tau_yz */
724:         {
725:           PetscScalar dx = 1.0, dy = 1.0, dz = 1.0; /* initialization to prevent incorrect compiler warnings */

727:           if (ex > 0) dx = arr_coord_x[ex][slot_coord_element] - arr_coord_x[ex - 1][slot_coord_element];
728:           if (ey > 0) dy = arr_coord_y[ey][slot_coord_element] - arr_coord_y[ey - 1][slot_coord_element];
729:           if (ez > 0) dz = arr_coord_z[ez][slot_coord_element] - arr_coord_z[ez - 1][slot_coord_element];

731:           if (ex > 0 && ey > 0) {
732:             const PetscScalar M = arr_lame[ez][ey][ex][slot_mu_downleft];

734:             arr_stress[ez][ey][ex][slot_txy_downleft] += M * ctx->dt * ((arr_velocity[ez][ey][ex][slot_vx_left] - arr_velocity[ez][ey - 1][ex][slot_vx_left]) / dy + (arr_velocity[ez][ey][ex][slot_vy_down] - arr_velocity[ez][ey][ex - 1][slot_vy_down]) / dx);
735:           }

737:           /* Update tau_xz */
738:           if (ex > 0 && ez > 0) {
739:             const PetscScalar M = arr_lame[ez][ey][ex][slot_mu_backleft];

741:             arr_stress[ez][ey][ex][slot_txz_backleft] += M * ctx->dt * ((arr_velocity[ez][ey][ex][slot_vx_left] - arr_velocity[ez - 1][ey][ex][slot_vx_left]) / dz + (arr_velocity[ez][ey][ex][slot_vz_back] - arr_velocity[ez][ey][ex - 1][slot_vz_back]) / dx);
742:           }

744:           /* Update tau_yz */
745:           if (ey > 0 && ez > 0) {
746:             const PetscScalar M = arr_lame[ez][ey][ex][slot_mu_backdown];

748:             arr_stress[ez][ey][ex][slot_tyz_backdown] += M * ctx->dt * ((arr_velocity[ez][ey][ex][slot_vy_down] - arr_velocity[ez - 1][ey][ex][slot_vy_down]) / dz + (arr_velocity[ez][ey][ex][slot_vz_back] - arr_velocity[ez][ey - 1][ex][slot_vz_back]) / dy);
749:           }
750:         }
751:       }
752:     }
753:   }

755:   /* Restore all access */
756:   DMStagVecRestoreArray(ctx->dm_stress, stress_local, &arr_stress);
757:   DMLocalToGlobal(ctx->dm_stress, stress_local, INSERT_VALUES, stress);
758:   DMRestoreLocalVector(ctx->dm_stress, &stress_local);
759:   DMStagVecRestoreArrayRead(ctx->dm_velocity, velocity_local, (void *)&arr_velocity);
760:   DMRestoreLocalVector(ctx->dm_velocity, &velocity_local);
761:   DMStagVecRestoreArrayRead(ctx->dm_lame, lame_local, (void *)&arr_lame);
762:   DMRestoreLocalVector(ctx->dm_lame, &lame_local);
763:   DMStagRestoreProductCoordinateArrays(ctx->dm_velocity, (void *)&arr_coord_x, (void *)&arr_coord_y, (void *)&arr_coord_z);
764:   return 0;
765: }

767: static PetscErrorCode UpdateStress(const Ctx *ctx, Vec velocity, Vec stress, Vec lame)
768: {
770:   if (ctx->dim == 2) {
771:     UpdateStress_2d(ctx, velocity, stress, lame);
772:   } else if (ctx->dim == 3) {
773:     UpdateStress_3d(ctx, velocity, stress, lame);
774:   }
775:   return 0;
776: }

778: static PetscErrorCode DumpStress(const Ctx *ctx, Vec stress, PetscInt timestep)
779: {
780:   DM  da_normal, da_shear   = NULL;
781:   Vec vec_normal, vec_shear = NULL;


785:   DMStagVecSplitToDMDA(ctx->dm_stress, stress, DMSTAG_ELEMENT, -ctx->dim, &da_normal, &vec_normal);
786:   PetscObjectSetName((PetscObject)vec_normal, "normal stresses");

788:   /* Dump element-based fields to a .vtr file */
789:   {
790:     PetscViewer viewer;
791:     char        filename[PETSC_MAX_PATH_LEN];

793:     PetscSNPrintf(filename, sizeof(filename), "ex6_stress_normal_%.4" PetscInt_FMT ".vtr", timestep);
794:     PetscViewerVTKOpen(PetscObjectComm((PetscObject)da_normal), filename, FILE_MODE_WRITE, &viewer);
795:     VecView(vec_normal, viewer);
796:     PetscViewerDestroy(&viewer);
797:     PetscPrintf(PETSC_COMM_WORLD, "Created %s\n", filename);
798:   }

800:   if (ctx->dim == 2) {
801:     /* (2D only) Dump vertex-based fields to a second .vtr file */
802:     DMStagVecSplitToDMDA(ctx->dm_stress, stress, DMSTAG_DOWN_LEFT, 0, &da_shear, &vec_shear);
803:     PetscObjectSetName((PetscObject)vec_shear, "shear stresses");

805:     {
806:       PetscViewer viewer;
807:       char        filename[PETSC_MAX_PATH_LEN];

809:       PetscSNPrintf(filename, sizeof(filename), "ex6_stress_shear_%.4" PetscInt_FMT ".vtr", timestep);
810:       PetscViewerVTKOpen(PetscObjectComm((PetscObject)da_normal), filename, FILE_MODE_WRITE, &viewer);
811:       VecView(vec_shear, viewer);
812:       PetscViewerDestroy(&viewer);
813:       PetscPrintf(PETSC_COMM_WORLD, "Created %s\n", filename);
814:     }
815:   }

817:   /* Destroy DMDAs and Vecs */
818:   DMDestroy(&da_normal);
819:   DMDestroy(&da_shear);
820:   VecDestroy(&vec_normal);
821:   VecDestroy(&vec_shear);
822:   return 0;
823: }

825: static PetscErrorCode DumpVelocity(const Ctx *ctx, Vec velocity, PetscInt timestep)
826: {
827:   DM       dmVelAvg;
828:   Vec      velAvg;
829:   DM       daVelAvg;
830:   Vec      vecVelAvg;
831:   Vec      velocity_local;
832:   PetscInt ex, ey, ez, startx, starty, startz, nx, ny, nz;


836:   if (ctx->dim == 2) {
837:     DMStagCreateCompatibleDMStag(ctx->dm_velocity, 0, 0, 2, 0, &dmVelAvg); /* 2 dof per element */
838:   } else if (ctx->dim == 3) {
839:     DMStagCreateCompatibleDMStag(ctx->dm_velocity, 0, 0, 0, 3, &dmVelAvg); /* 3 dof per element */
840:   } else SETERRQ(PetscObjectComm((PetscObject)ctx->dm_velocity), PETSC_ERR_SUP, "Unsupported dim %" PetscInt_FMT, ctx->dim);
841:   DMSetUp(dmVelAvg);
842:   DMStagSetUniformCoordinatesProduct(dmVelAvg, ctx->xmin, ctx->xmax, ctx->ymin, ctx->ymax, ctx->zmin, ctx->zmax);
843:   DMCreateGlobalVector(dmVelAvg, &velAvg);
844:   DMGetLocalVector(ctx->dm_velocity, &velocity_local);
845:   DMGlobalToLocal(ctx->dm_velocity, velocity, INSERT_VALUES, velocity_local);
846:   DMStagGetCorners(dmVelAvg, &startx, &starty, &startz, &nx, &ny, &nz, NULL, NULL, NULL);
847:   if (ctx->dim == 2) {
848:     for (ey = starty; ey < starty + ny; ++ey) {
849:       for (ex = startx; ex < startx + nx; ++ex) {
850:         DMStagStencil from[4], to[2];
851:         PetscScalar   valFrom[4], valTo[2];

853:         from[0].i   = ex;
854:         from[0].j   = ey;
855:         from[0].loc = DMSTAG_UP;
856:         from[0].c   = 0;
857:         from[1].i   = ex;
858:         from[1].j   = ey;
859:         from[1].loc = DMSTAG_DOWN;
860:         from[1].c   = 0;
861:         from[2].i   = ex;
862:         from[2].j   = ey;
863:         from[2].loc = DMSTAG_LEFT;
864:         from[2].c   = 0;
865:         from[3].i   = ex;
866:         from[3].j   = ey;
867:         from[3].loc = DMSTAG_RIGHT;
868:         from[3].c   = 0;
869:         DMStagVecGetValuesStencil(ctx->dm_velocity, velocity_local, 4, from, valFrom);
870:         to[0].i   = ex;
871:         to[0].j   = ey;
872:         to[0].loc = DMSTAG_ELEMENT;
873:         to[0].c   = 0;
874:         valTo[0]  = 0.5 * (valFrom[2] + valFrom[3]);
875:         to[1].i   = ex;
876:         to[1].j   = ey;
877:         to[1].loc = DMSTAG_ELEMENT;
878:         to[1].c   = 1;
879:         valTo[1]  = 0.5 * (valFrom[0] + valFrom[1]);
880:         DMStagVecSetValuesStencil(dmVelAvg, velAvg, 2, to, valTo, INSERT_VALUES);
881:       }
882:     }
883:   } else if (ctx->dim == 3) {
884:     for (ez = startz; ez < startz + nz; ++ez) {
885:       for (ey = starty; ey < starty + ny; ++ey) {
886:         for (ex = startx; ex < startx + nx; ++ex) {
887:           DMStagStencil from[6], to[3];
888:           PetscScalar   valFrom[6], valTo[3];

890:           from[0].i   = ex;
891:           from[0].j   = ey;
892:           from[0].k   = ez;
893:           from[0].loc = DMSTAG_UP;
894:           from[0].c   = 0;
895:           from[1].i   = ex;
896:           from[1].j   = ey;
897:           from[1].k   = ez;
898:           from[1].loc = DMSTAG_DOWN;
899:           from[1].c   = 0;
900:           from[2].i   = ex;
901:           from[2].j   = ey;
902:           from[2].k   = ez;
903:           from[2].loc = DMSTAG_LEFT;
904:           from[2].c   = 0;
905:           from[3].i   = ex;
906:           from[3].j   = ey;
907:           from[3].k   = ez;
908:           from[3].loc = DMSTAG_RIGHT;
909:           from[3].c   = 0;
910:           from[4].i   = ex;
911:           from[4].j   = ey;
912:           from[4].k   = ez;
913:           from[4].loc = DMSTAG_FRONT;
914:           from[4].c   = 0;
915:           from[5].i   = ex;
916:           from[5].j   = ey;
917:           from[5].k   = ez;
918:           from[5].loc = DMSTAG_BACK;
919:           from[5].c   = 0;
920:           DMStagVecGetValuesStencil(ctx->dm_velocity, velocity_local, 6, from, valFrom);
921:           to[0].i   = ex;
922:           to[0].j   = ey;
923:           to[0].k   = ez;
924:           to[0].loc = DMSTAG_ELEMENT;
925:           to[0].c   = 0;
926:           valTo[0]  = 0.5 * (valFrom[2] + valFrom[3]);
927:           to[1].i   = ex;
928:           to[1].j   = ey;
929:           to[1].k   = ez;
930:           to[1].loc = DMSTAG_ELEMENT;
931:           to[1].c   = 1;
932:           valTo[1]  = 0.5 * (valFrom[0] + valFrom[1]);
933:           to[2].i   = ex;
934:           to[2].j   = ey;
935:           to[2].k   = ez;
936:           to[2].loc = DMSTAG_ELEMENT;
937:           to[2].c   = 2;
938:           valTo[2]  = 0.5 * (valFrom[4] + valFrom[5]);
939:           DMStagVecSetValuesStencil(dmVelAvg, velAvg, 3, to, valTo, INSERT_VALUES);
940:         }
941:       }
942:     }
943:   } else SETERRQ(PetscObjectComm((PetscObject)ctx->dm_velocity), PETSC_ERR_SUP, "Unsupported dim %" PetscInt_FMT, ctx->dim);
944:   VecAssemblyBegin(velAvg);
945:   VecAssemblyEnd(velAvg);
946:   DMRestoreLocalVector(ctx->dm_velocity, &velocity_local);

948:   DMStagVecSplitToDMDA(dmVelAvg, velAvg, DMSTAG_ELEMENT, -3, &daVelAvg, &vecVelAvg); /* note -3 : pad with zero in 2D case */
949:   PetscObjectSetName((PetscObject)vecVelAvg, "Velocity (Averaged)");

951:   /* Dump element-based fields to a .vtr file */
952:   {
953:     PetscViewer viewer;
954:     char        filename[PETSC_MAX_PATH_LEN];

956:     PetscSNPrintf(filename, sizeof(filename), "ex6_velavg_%.4" PetscInt_FMT ".vtr", timestep);
957:     PetscViewerVTKOpen(PetscObjectComm((PetscObject)daVelAvg), filename, FILE_MODE_WRITE, &viewer);
958:     VecView(vecVelAvg, viewer);
959:     PetscViewerDestroy(&viewer);
960:     PetscPrintf(PETSC_COMM_WORLD, "Created %s\n", filename);
961:   }

963:   /* Destroy DMDAs and Vecs */
964:   VecDestroy(&vecVelAvg);
965:   DMDestroy(&daVelAvg);
966:   VecDestroy(&velAvg);
967:   DMDestroy(&dmVelAvg);
968:   return 0;
969: }

971: /*TEST

973:    test:
974:       suffix: 1
975:       requires: !complex
976:       nsize: 1
977:       args: -stag_grid_x 12 -stag_grid_y 7 -nsteps 4 -dump_output 0

979:    test:
980:       suffix: 2
981:       requires: !complex
982:       nsize: 9
983:       args: -stag_grid_x 12 -stag_grid_y 15 -nsteps 3 -dump_output 0

985:    test:
986:       suffix: 3
987:       requires: !complex
988:       nsize: 1
989:       args: -stag_grid_x 12 -stag_grid_y 7 -stag_grid_z 11 -nsteps 2 -dim 3 -dump_output 0

991:    test:
992:       suffix: 4
993:       requires: !complex
994:       nsize: 12
995:       args: -stag_grid_x 12 -stag_grid_y 15 -stag_grid_z 8 -nsteps 3 -dim 3 -dump_output 0

997: TEST*/