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;

 64:   PetscFunctionBeginUser;
 65:   /* Initialize PETSc */
 66:   PetscCall(PetscInitialize(&argc, &argv, 0, help));

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

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

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

 94:     switch (ctx.dim) {
 95:     case 2:
 96:       dof0 = 0;
 97:       dof1 = 1;
 98:       dof2 = 0; /* 1 dof per cell boundary */
 99:       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));
100:       break;
101:     case 3:
102:       dof0 = 0;
103:       dof1 = 0;
104:       dof2 = 1;
105:       dof3 = 0; /* 1 dof per cell boundary */
106:       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));
107:       break;
108:     default:
109:       SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_SUP, "Not Implemented for dimension %" PetscInt_FMT, ctx.dim);
110:     }
111:   }
112:   PetscCall(DMSetFromOptions(ctx.dm_velocity)); /* Options control velocity DM */
113:   PetscCall(DMSetUp(ctx.dm_velocity));
114:   PetscCall(DMStagSetUniformCoordinatesProduct(ctx.dm_velocity, ctx.xmin, ctx.xmax, ctx.ymin, ctx.ymax, ctx.zmin, ctx.zmax));

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

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

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

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

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

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

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

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

196:   /* Time Loop */
197:   for (PetscInt timestep = 1; timestep <= ctx.timesteps; ++timestep) {
198:     const PetscReal t = timestep * ctx.dt;

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

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

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

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

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

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

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

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

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

329:   PetscFunctionBeginUser;
330:   PetscCall(DMStagGetGlobalSizes(ctx->dm_stress, &N[0], &N[1], &N[2]));
331:   PetscCall(DMStagGetCorners(ctx->dm_stress, &start[0], &start[1], &start[2], &n[0], &n[1], &n[2], NULL, NULL, NULL));

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

361:   PetscCall(VecAssemblyBegin(stress));
362:   PetscCall(VecAssemblyEnd(stress));
363:   PetscFunctionReturn(PETSC_SUCCESS);
364: }

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

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

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

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

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

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

418:         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);
419:       }

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

427:         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);
428:       }
429:     }
430:   }

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

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

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

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

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

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

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

510:           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);
511:         }

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

520:           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);
521:         }

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

530:           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);
531:         }
532:       }
533:     }
534:   }

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

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

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

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

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

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

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

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

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

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

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

625:         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);
626:       }
627:     }
628:   }

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

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

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

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

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

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

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

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

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

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

716:         /* Update tau_xy, tau_xz, tau_yz */
717:         {
718:           PetscScalar dx = 1.0, dy = 1.0, dz = 1.0; /* initialization to prevent incorrect compiler warnings */

720:           if (ex > 0) dx = arr_coord_x[ex][slot_coord_element] - arr_coord_x[ex - 1][slot_coord_element];
721:           if (ey > 0) dy = arr_coord_y[ey][slot_coord_element] - arr_coord_y[ey - 1][slot_coord_element];
722:           if (ez > 0) dz = arr_coord_z[ez][slot_coord_element] - arr_coord_z[ez - 1][slot_coord_element];

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

727:             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);
728:           }

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

734:             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);
735:           }

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

741:             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);
742:           }
743:         }
744:       }
745:     }
746:   }

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

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

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

776:   PetscFunctionBeginUser;
777:   PetscCall(DMStagVecSplitToDMDA(ctx->dm_stress, stress, DMSTAG_ELEMENT, -ctx->dim, &da_normal, &vec_normal));
778:   PetscCall(PetscObjectSetName((PetscObject)vec_normal, "normal stresses"));

780:   /* Dump element-based fields to a .vtr file */
781:   {
782:     PetscViewer viewer;
783:     char        filename[PETSC_MAX_PATH_LEN];

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

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

797:     {
798:       PetscViewer viewer;
799:       char        filename[PETSC_MAX_PATH_LEN];

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

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

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

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

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

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

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

942:   /* Dump element-based fields to a .vtr file */
943:   {
944:     PetscViewer viewer;
945:     char        filename[PETSC_MAX_PATH_LEN];

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

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

962: /*TEST

964:    test:
965:       suffix: 1
966:       requires: !complex
967:       nsize: 1
968:       args: -stag_grid_x 12 -stag_grid_y 7 -nsteps 4 -dump_output 0

970:    test:
971:       suffix: 2
972:       requires: !complex
973:       nsize: 9
974:       args: -stag_grid_x 12 -stag_grid_y 15 -nsteps 3 -dump_output 0

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

982:    test:
983:       suffix: 4
984:       requires: !complex
985:       nsize: 12
986:       args: -stag_grid_x 12 -stag_grid_y 15 -stag_grid_z 8 -nsteps 3 -dim 3 -dump_output 0

988: TEST*/