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;

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

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

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

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

412:   /* Iterate over interior of the domain, updating the velocities */
413:   PetscCall(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:   PetscCall(DMStagVecRestoreArrayRead(ctx->dm_buoyancy, buoyancy_local, (void *)&arr_buoyancy));
438:   PetscCall(DMRestoreLocalVector(ctx->dm_buoyancy, &buoyancy_local));
439:   PetscCall(DMStagVecRestoreArrayRead(ctx->dm_stress, stress_local, (void *)&arr_stress));
440:   PetscCall(DMRestoreLocalVector(ctx->dm_stress, &stress_local));
441:   PetscCall(DMStagVecRestoreArray(ctx->dm_velocity, velocity_local, &arr_velocity));
442:   PetscCall(DMLocalToGlobal(ctx->dm_velocity, velocity_local, INSERT_VALUES, velocity));
443:   PetscCall(DMRestoreLocalVector(ctx->dm_velocity, &velocity_local));
444:   PetscCall(DMStagRestoreProductCoordinateArrays(ctx->dm_velocity, (void *)&arr_coord_x, (void *)&arr_coord_y, NULL));
445:   PetscFunctionReturn(PETSC_SUCCESS);
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;

462:   PetscFunctionBeginUser;

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

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

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

497:   /* Prepare read-only access to coordinate data */
498:   PetscCall(DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_LEFT, &slot_coord_prev));
499:   PetscCall(DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_RIGHT, &slot_coord_next));
500:   PetscCall(DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_ELEMENT, &slot_coord_element));
501:   PetscCall(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:   PetscCall(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:   PetscCall(DMStagVecRestoreArrayRead(ctx->dm_buoyancy, buoyancy_local, (void *)&arr_buoyancy));
543:   PetscCall(DMRestoreLocalVector(ctx->dm_buoyancy, &buoyancy_local));
544:   PetscCall(DMStagVecRestoreArrayRead(ctx->dm_stress, stress_local, (void *)&arr_stress));
545:   PetscCall(DMRestoreLocalVector(ctx->dm_stress, &stress_local));
546:   PetscCall(DMStagVecRestoreArray(ctx->dm_velocity, velocity_local, &arr_velocity));
547:   PetscCall(DMLocalToGlobal(ctx->dm_velocity, velocity_local, INSERT_VALUES, velocity));
548:   PetscCall(DMRestoreLocalVector(ctx->dm_velocity, &velocity_local));
549:   PetscCall(DMStagRestoreProductCoordinateArrays(ctx->dm_velocity, (void *)&arr_coord_x, (void *)&arr_coord_y, (void *)&arr_coord_z));
550:   PetscFunctionReturn(PETSC_SUCCESS);
551: }

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

576:   PetscFunctionBeginUser;

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

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

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

603:   /* Prepare read-only access to coordinate data */
604:   PetscCall(DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_LEFT, &slot_coord_prev));
605:   PetscCall(DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_RIGHT, &slot_coord_next));
606:   PetscCall(DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_ELEMENT, &slot_coord_element));
607:   PetscCall(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:   PetscCall(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:   PetscCall(DMStagVecRestoreArray(ctx->dm_stress, stress_local, &arr_stress));
638:   PetscCall(DMLocalToGlobal(ctx->dm_stress, stress_local, INSERT_VALUES, stress));
639:   PetscCall(DMRestoreLocalVector(ctx->dm_stress, &stress_local));
640:   PetscCall(DMStagVecRestoreArrayRead(ctx->dm_velocity, velocity_local, (void *)&arr_velocity));
641:   PetscCall(DMRestoreLocalVector(ctx->dm_velocity, &velocity_local));
642:   PetscCall(DMStagVecRestoreArrayRead(ctx->dm_lame, lame_local, (void *)&arr_lame));
643:   PetscCall(DMRestoreLocalVector(ctx->dm_lame, &lame_local));
644:   PetscCall(DMStagRestoreProductCoordinateArrays(ctx->dm_velocity, (void *)&arr_coord_x, (void *)&arr_coord_y, NULL));
645:   PetscFunctionReturn(PETSC_SUCCESS);
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;

660:   PetscFunctionBeginUser;

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

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

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

694:   /* Prepare read-only access to coordinate data */
695:   PetscCall(DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_LEFT, &slot_coord_prev));
696:   PetscCall(DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_RIGHT, &slot_coord_next));
697:   PetscCall(DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_ELEMENT, &slot_coord_element));
698:   PetscCall(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:   PetscCall(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:   PetscCall(DMStagVecRestoreArray(ctx->dm_stress, stress_local, &arr_stress));
757:   PetscCall(DMLocalToGlobal(ctx->dm_stress, stress_local, INSERT_VALUES, stress));
758:   PetscCall(DMRestoreLocalVector(ctx->dm_stress, &stress_local));
759:   PetscCall(DMStagVecRestoreArrayRead(ctx->dm_velocity, velocity_local, (void *)&arr_velocity));
760:   PetscCall(DMRestoreLocalVector(ctx->dm_velocity, &velocity_local));
761:   PetscCall(DMStagVecRestoreArrayRead(ctx->dm_lame, lame_local, (void *)&arr_lame));
762:   PetscCall(DMRestoreLocalVector(ctx->dm_lame, &lame_local));
763:   PetscCall(DMStagRestoreProductCoordinateArrays(ctx->dm_velocity, (void *)&arr_coord_x, (void *)&arr_coord_y, (void *)&arr_coord_z));
764:   PetscFunctionReturn(PETSC_SUCCESS);
765: }

767: static PetscErrorCode UpdateStress(const Ctx *ctx, Vec velocity, Vec stress, Vec lame)
768: {
769:   PetscFunctionBeginUser;
770:   if (ctx->dim == 2) {
771:     PetscCall(UpdateStress_2d(ctx, velocity, stress, lame));
772:   } else if (ctx->dim == 3) {
773:     PetscCall(UpdateStress_3d(ctx, velocity, stress, lame));
774:   }
775:   PetscFunctionReturn(PETSC_SUCCESS);
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;

783:   PetscFunctionBeginUser;

785:   PetscCall(DMStagVecSplitToDMDA(ctx->dm_stress, stress, DMSTAG_ELEMENT, -ctx->dim, &da_normal, &vec_normal));
786:   PetscCall(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:     PetscCall(PetscSNPrintf(filename, sizeof(filename), "ex6_stress_normal_%.4" PetscInt_FMT ".vtr", timestep));
794:     PetscCall(PetscViewerVTKOpen(PetscObjectComm((PetscObject)da_normal), filename, FILE_MODE_WRITE, &viewer));
795:     PetscCall(VecView(vec_normal, viewer));
796:     PetscCall(PetscViewerDestroy(&viewer));
797:     PetscCall(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:     PetscCall(DMStagVecSplitToDMDA(ctx->dm_stress, stress, DMSTAG_DOWN_LEFT, 0, &da_shear, &vec_shear));
803:     PetscCall(PetscObjectSetName((PetscObject)vec_shear, "shear stresses"));

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

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

817:   /* Destroy DMDAs and Vecs */
818:   PetscCall(DMDestroy(&da_normal));
819:   PetscCall(DMDestroy(&da_shear));
820:   PetscCall(VecDestroy(&vec_normal));
821:   PetscCall(VecDestroy(&vec_shear));
822:   PetscFunctionReturn(PETSC_SUCCESS);
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;

834:   PetscFunctionBeginUser;

836:   if (ctx->dim == 2) {
837:     PetscCall(DMStagCreateCompatibleDMStag(ctx->dm_velocity, 0, 0, 2, 0, &dmVelAvg)); /* 2 dof per element */
838:   } else if (ctx->dim == 3) {
839:     PetscCall(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:   PetscCall(DMSetUp(dmVelAvg));
842:   PetscCall(DMStagSetUniformCoordinatesProduct(dmVelAvg, ctx->xmin, ctx->xmax, ctx->ymin, ctx->ymax, ctx->zmin, ctx->zmax));
843:   PetscCall(DMCreateGlobalVector(dmVelAvg, &velAvg));
844:   PetscCall(DMGetLocalVector(ctx->dm_velocity, &velocity_local));
845:   PetscCall(DMGlobalToLocal(ctx->dm_velocity, velocity, INSERT_VALUES, velocity_local));
846:   PetscCall(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:         PetscCall(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:         PetscCall(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:           PetscCall(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:           PetscCall(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:   PetscCall(VecAssemblyBegin(velAvg));
945:   PetscCall(VecAssemblyEnd(velAvg));
946:   PetscCall(DMRestoreLocalVector(ctx->dm_velocity, &velocity_local));

948:   PetscCall(DMStagVecSplitToDMDA(dmVelAvg, velAvg, DMSTAG_ELEMENT, -3, &daVelAvg, &vecVelAvg)); /* note -3 : pad with zero in 2D case */
949:   PetscCall(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:     PetscCall(PetscSNPrintf(filename, sizeof(filename), "ex6_velavg_%.4" PetscInt_FMT ".vtr", timestep));
957:     PetscCall(PetscViewerVTKOpen(PetscObjectComm((PetscObject)daVelAvg), filename, FILE_MODE_WRITE, &viewer));
958:     PetscCall(VecView(vecVelAvg, viewer));
959:     PetscCall(PetscViewerDestroy(&viewer));
960:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Created %s\n", filename));
961:   }

963:   /* Destroy DMDAs and Vecs */
964:   PetscCall(VecDestroy(&vecVelAvg));
965:   PetscCall(DMDestroy(&daVelAvg));
966:   PetscCall(VecDestroy(&velAvg));
967:   PetscCall(DMDestroy(&dmVelAvg));
968:   PetscFunctionReturn(PETSC_SUCCESS);
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*/