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 */
 66:   PetscFunctionBeginUser;
 67:   PetscCall(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:   PetscCall(PetscOptionsGetInt(NULL, NULL, "-dim", &ctx.dim, NULL));
 86:   PetscCall(PetscOptionsGetReal(NULL, NULL, "-dt", &ctx.dt, NULL));
 87:   PetscCall(PetscOptionsGetInt(NULL, NULL, "-nsteps", &ctx.timesteps, NULL));
 88:   PetscCall(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:       PetscCall(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:       PetscCall(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:   PetscCall(DMSetFromOptions(ctx.dm_velocity)); /* Options control velocity DM */
114:   PetscCall(DMSetUp(ctx.dm_velocity));
115:   PetscCall(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:     PetscCall(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:     PetscCall(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:   PetscCall(DMSetUp(ctx.dm_stress));
131:   PetscCall(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:     PetscCall(DMStagCreateCompatibleDMStag(ctx.dm_velocity, 0, 1, 0, 0, &ctx.dm_buoyancy));
138:     break;
139:   case 3:
140:     /* buoyancy on element boundaries (faces) */
141:     PetscCall(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:   PetscCall(DMSetUp(ctx.dm_buoyancy));

148:   switch (ctx.dim) {
149:   case 2:
150:     /* mu and lambda + 2*mu on element centers, mu on corners */
151:     PetscCall(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:     PetscCall(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:   PetscCall(DMSetUp(ctx.dm_lame));

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

167:     PetscCall(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:       PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Using a %" PetscInt_FMT " x %" PetscInt_FMT " mesh\n", N[0], N[1]));
172:     } else if (ctx.dim == 3) {
173:       PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Using a %" PetscInt_FMT " x %" PetscInt_FMT " x %" PetscInt_FMT " mesh\n", N[0], N[1], N[2]));
174:     }
175:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "dx: %g\n", (double)PetscRealPart(dx)));
176:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "dt: %g\n", (double)ctx.dt));
177:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "P-wave velocity: %g\n", (double)PetscRealPart(PetscSqrtScalar((ctx.lambda + 2 * ctx.mu) / ctx.rho))));
178:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "V_p dt / dx: %g\n", (double)PetscRealPart(Vp * ctx.dt / dx)));
179:   }

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

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

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

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

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

213:   /* Clean up and finalize PETSc */
214:   PetscCall(VecDestroy(&velocity));
215:   PetscCall(VecDestroy(&stress));
216:   PetscCall(VecDestroy(&ctx.lame));
217:   PetscCall(VecDestroy(&ctx.buoyancy));
218:   PetscCall(DMDestroy(&ctx.dm_velocity));
219:   PetscCall(DMDestroy(&ctx.dm_stress));
220:   PetscCall(DMDestroy(&ctx.dm_buoyancy));
221:   PetscCall(DMDestroy(&ctx.dm_lame));
222:   PetscCall(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;

230:   PetscFunctionBeginUser;
231:   PetscCall(DMCreateGlobalVector(ctx->dm_lame, &ctx->lame));
232:   PetscCall(DMStagGetGlobalSizes(ctx->dm_lame, &N[0], &N[1], &N[2]));
233:   PetscCall(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:         PetscCall(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:         PetscCall(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:         PetscCall(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:           PetscCall(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:           PetscCall(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:             PetscCall(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:             PetscCall(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:             PetscCall(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:   PetscCall(VecAssemblyBegin(ctx->lame));
321:   PetscCall(VecAssemblyEnd(ctx->lame));
322:   PetscFunctionReturn(PETSC_SUCCESS);
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);

332:   PetscFunctionBeginUser;
333:   PetscCall(DMStagGetGlobalSizes(ctx->dm_stress, &N[0], &N[1], &N[2]));
334:   PetscCall(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:     PetscCall(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:     PetscCall(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:       PetscCall(DMStagVecSetValuesStencil(ctx->dm_stress, stress, 1, &pos, &val, INSERT_VALUES));
361:     }
362:   }

364:   PetscCall(VecAssemblyBegin(stress));
365:   PetscCall(VecAssemblyEnd(stress));
366:   PetscFunctionReturn(PETSC_SUCCESS);
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;

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

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

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

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

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

421:         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);
422:       }

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

430:         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);
431:       }
432:     }
433:   }

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

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

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

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

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

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

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

513:           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);
514:         }

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

523:           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);
524:         }

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

533:           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);
534:         }
535:       }
536:     }
537:   }

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

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

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

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

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

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

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

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

617:         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;

619:         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;
620:       }

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

628:         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);
629:       }
630:     }
631:   }

633:   /* Restore all access */
634:   PetscCall(DMStagVecRestoreArray(ctx->dm_stress, stress_local, &arr_stress));
635:   PetscCall(DMLocalToGlobal(ctx->dm_stress, stress_local, INSERT_VALUES, stress));
636:   PetscCall(DMRestoreLocalVector(ctx->dm_stress, &stress_local));
637:   PetscCall(DMStagVecRestoreArrayRead(ctx->dm_velocity, velocity_local, (void *)&arr_velocity));
638:   PetscCall(DMRestoreLocalVector(ctx->dm_velocity, &velocity_local));
639:   PetscCall(DMStagVecRestoreArrayRead(ctx->dm_lame, lame_local, (void *)&arr_lame));
640:   PetscCall(DMRestoreLocalVector(ctx->dm_lame, &lame_local));
641:   PetscCall(DMStagRestoreProductCoordinateArrays(ctx->dm_velocity, (void *)&arr_coord_x, (void *)&arr_coord_y, NULL));
642:   PetscFunctionReturn(PETSC_SUCCESS);
643: }

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

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

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

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

690:   /* Prepare read-only access to coordinate data */
691:   PetscCall(DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_LEFT, &slot_coord_prev));
692:   PetscCall(DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_RIGHT, &slot_coord_next));
693:   PetscCall(DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_ELEMENT, &slot_coord_element));
694:   PetscCall(DMStagGetProductCoordinateArrays(ctx->dm_velocity, (void *)&arr_coord_x, (void *)&arr_coord_y, (void *)&arr_coord_z));

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

709:           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 +
710:                                               L * ctx->dt * (arr_velocity[ez][ey][ex][slot_vz_front] - arr_velocity[ez][ey][ex][slot_vz_back]) / dz;

712:           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 +
713:                                               L * ctx->dt * (arr_velocity[ez][ey][ex][slot_vx_right] - arr_velocity[ez][ey][ex][slot_vx_left]) / dx;

715:           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 +
716:                                               L * ctx->dt * (arr_velocity[ez][ey][ex][slot_vy_up] - arr_velocity[ez][ey][ex][slot_vy_down]) / dy;
717:         }

719:         /* Update tau_xy, tau_xz, tau_yz */
720:         {
721:           PetscScalar dx = 1.0, dy = 1.0, dz = 1.0; /* initialization to prevent incorrect compiler warnings */

723:           if (ex > 0) dx = arr_coord_x[ex][slot_coord_element] - arr_coord_x[ex - 1][slot_coord_element];
724:           if (ey > 0) dy = arr_coord_y[ey][slot_coord_element] - arr_coord_y[ey - 1][slot_coord_element];
725:           if (ez > 0) dz = arr_coord_z[ez][slot_coord_element] - arr_coord_z[ez - 1][slot_coord_element];

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

730:             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);
731:           }

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

737:             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);
738:           }

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

744:             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);
745:           }
746:         }
747:       }
748:     }
749:   }

751:   /* Restore all access */
752:   PetscCall(DMStagVecRestoreArray(ctx->dm_stress, stress_local, &arr_stress));
753:   PetscCall(DMLocalToGlobal(ctx->dm_stress, stress_local, INSERT_VALUES, stress));
754:   PetscCall(DMRestoreLocalVector(ctx->dm_stress, &stress_local));
755:   PetscCall(DMStagVecRestoreArrayRead(ctx->dm_velocity, velocity_local, (void *)&arr_velocity));
756:   PetscCall(DMRestoreLocalVector(ctx->dm_velocity, &velocity_local));
757:   PetscCall(DMStagVecRestoreArrayRead(ctx->dm_lame, lame_local, (void *)&arr_lame));
758:   PetscCall(DMRestoreLocalVector(ctx->dm_lame, &lame_local));
759:   PetscCall(DMStagRestoreProductCoordinateArrays(ctx->dm_velocity, (void *)&arr_coord_x, (void *)&arr_coord_y, (void *)&arr_coord_z));
760:   PetscFunctionReturn(PETSC_SUCCESS);
761: }

763: static PetscErrorCode UpdateStress(const Ctx *ctx, Vec velocity, Vec stress, Vec lame)
764: {
765:   PetscFunctionBeginUser;
766:   if (ctx->dim == 2) {
767:     PetscCall(UpdateStress_2d(ctx, velocity, stress, lame));
768:   } else if (ctx->dim == 3) {
769:     PetscCall(UpdateStress_3d(ctx, velocity, stress, lame));
770:   }
771:   PetscFunctionReturn(PETSC_SUCCESS);
772: }

774: static PetscErrorCode DumpStress(const Ctx *ctx, Vec stress, PetscInt timestep)
775: {
776:   DM  da_normal, da_shear   = NULL;
777:   Vec vec_normal, vec_shear = NULL;

779:   PetscFunctionBeginUser;
780:   PetscCall(DMStagVecSplitToDMDA(ctx->dm_stress, stress, DMSTAG_ELEMENT, -ctx->dim, &da_normal, &vec_normal));
781:   PetscCall(PetscObjectSetName((PetscObject)vec_normal, "normal stresses"));

783:   /* Dump element-based fields to a .vtr file */
784:   {
785:     PetscViewer viewer;
786:     char        filename[PETSC_MAX_PATH_LEN];

788:     PetscCall(PetscSNPrintf(filename, sizeof(filename), "ex6_stress_normal_%.4" PetscInt_FMT ".vtr", timestep));
789:     PetscCall(PetscViewerVTKOpen(PetscObjectComm((PetscObject)da_normal), filename, FILE_MODE_WRITE, &viewer));
790:     PetscCall(VecView(vec_normal, viewer));
791:     PetscCall(PetscViewerDestroy(&viewer));
792:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Created %s\n", filename));
793:   }

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

800:     {
801:       PetscViewer viewer;
802:       char        filename[PETSC_MAX_PATH_LEN];

804:       PetscCall(PetscSNPrintf(filename, sizeof(filename), "ex6_stress_shear_%.4" PetscInt_FMT ".vtr", timestep));
805:       PetscCall(PetscViewerVTKOpen(PetscObjectComm((PetscObject)da_normal), filename, FILE_MODE_WRITE, &viewer));
806:       PetscCall(VecView(vec_shear, viewer));
807:       PetscCall(PetscViewerDestroy(&viewer));
808:       PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Created %s\n", filename));
809:     }
810:   }

812:   /* Destroy DMDAs and Vecs */
813:   PetscCall(DMDestroy(&da_normal));
814:   PetscCall(DMDestroy(&da_shear));
815:   PetscCall(VecDestroy(&vec_normal));
816:   PetscCall(VecDestroy(&vec_shear));
817:   PetscFunctionReturn(PETSC_SUCCESS);
818: }

820: static PetscErrorCode DumpVelocity(const Ctx *ctx, Vec velocity, PetscInt timestep)
821: {
822:   DM       dmVelAvg;
823:   Vec      velAvg;
824:   DM       daVelAvg;
825:   Vec      vecVelAvg;
826:   Vec      velocity_local;
827:   PetscInt ex, ey, ez, startx, starty, startz, nx, ny, nz;

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

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

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

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

945:   /* Dump element-based fields to a .vtr file */
946:   {
947:     PetscViewer viewer;
948:     char        filename[PETSC_MAX_PATH_LEN];

950:     PetscCall(PetscSNPrintf(filename, sizeof(filename), "ex6_velavg_%.4" PetscInt_FMT ".vtr", timestep));
951:     PetscCall(PetscViewerVTKOpen(PetscObjectComm((PetscObject)daVelAvg), filename, FILE_MODE_WRITE, &viewer));
952:     PetscCall(VecView(vecVelAvg, viewer));
953:     PetscCall(PetscViewerDestroy(&viewer));
954:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Created %s\n", filename));
955:   }

957:   /* Destroy DMDAs and Vecs */
958:   PetscCall(VecDestroy(&vecVelAvg));
959:   PetscCall(DMDestroy(&daVelAvg));
960:   PetscCall(VecDestroy(&velAvg));
961:   PetscCall(DMDestroy(&dmVelAvg));
962:   PetscFunctionReturn(PETSC_SUCCESS);
963: }

965: /*TEST

967:    test:
968:       suffix: 1
969:       requires: !complex
970:       nsize: 1
971:       args: -stag_grid_x 12 -stag_grid_y 7 -nsteps 4 -dump_output 0

973:    test:
974:       suffix: 2
975:       requires: !complex
976:       nsize: 9
977:       args: -stag_grid_x 12 -stag_grid_y 15 -nsteps 3 -dump_output 0

979:    test:
980:       suffix: 3
981:       requires: !complex
982:       nsize: 1
983:       args: -stag_grid_x 12 -stag_grid_y 7 -stag_grid_z 11 -nsteps 2 -dim 3 -dump_output 0

985:    test:
986:       suffix: 4
987:       requires: !complex
988:       nsize: 12
989:       args: -stag_grid_x 12 -stag_grid_y 15 -stag_grid_z 8 -nsteps 3 -dim 3 -dump_output 0

991: TEST*/