Actual source code: ex6.c

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

      8: /*
      9: Reference:

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    364:   PetscCall(VecAssemblyBegin(stress));
    365:   PetscCall(VecAssemblyEnd(stress));
    366:   PetscFunctionReturn(PETSC_SUCCESS);
    367: }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    744:             arr_stress[ez][ey][ex][slot_tyz_backdown] += M * ctx->dt * ((arr_velocity[ez][ey][ex][slot_vy_down] - arr_velocity[ez - 1][ey][ex][slot_vy_down]) / dz + (arr_velocity[ez][ey][ex][slot_vz_back] - arr_velocity[ez][ey - 1][ex][slot_vz_back]) / dy);
    745:           }
    746:         }
    747:       }
    748:     }
    749:   }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    965: /*TEST

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

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

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

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

    991: TEST*/