Date: 2022-03-04T17:38:36+00:00

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: Notes:

     24:   This example uses y in 2 dimensions, where the paper uses z.
     25:   This example uses the dual grid of the one pictured in Fig. 1. of the paper,
     26: so that velocities are on face boundaries, shear stresses are defined on vertices,
     27: 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).
     31:   This example treats the boundaries naively (by leaving ~zero velocity and stress there).

     33: */

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

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

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

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

     65:   /* Initialize PETSc */
     66:   PetscInitialize(&argc,&argv,0,help);if (ierr) return ierr;

     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; ctx.xmax = 1.0;
     74:   ctx.ymin        = 0.0; ctx.ymax = 1.0;
     75:   ctx.zmin        = 0.0; ctx.zmax = 1.0;
     76:   ctx.dt          = 0.001;
     77:   ctx.timesteps   = 100;
     78:   ctx.dump_output = PETSC_TRUE;

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

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

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

    107:   /* Create a second, compatible DMStag for the stresses */
    108:   switch (ctx.dim) {
    109:     case 2:
    110:       /* One shear stress component on element corners, two shear stress components on elements */
    111:       DMStagCreateCompatibleDMStag(ctx.dm_velocity,1,0,2,0,&ctx.dm_stress);
    112:       break;
    113:     case 3:
    114:       /* One shear stress component on element edges, three shear stress components on elements */
    115:       DMStagCreateCompatibleDMStag(ctx.dm_velocity,0,1,0,3,&ctx.dm_stress);
    116:       break;
    117:     default: SETERRQ1(PETSC_COMM_WORLD,PETSC_ERR_SUP,"Not Implemented for dimension %D",ctx.dim);
    118:   }
    119:   DMSetUp(ctx.dm_stress);
    120:   DMStagSetUniformCoordinatesProduct(ctx.dm_stress,ctx.xmin,ctx.xmax,ctx.ymin,ctx.ymax,ctx.zmin,ctx.zmax);

    122:   /* Create two additional DMStag objects for the buoyancy and Lame parameters */
    123:   switch (ctx.dim) {
    124:     case 2:
    125:       /* buoyancy on element boundaries (edges) */
    126:       DMStagCreateCompatibleDMStag(ctx.dm_velocity,0,1,0,0,&ctx.dm_buoyancy);
    127:       break;
    128:     case 3:
    129:       /* buoyancy on element boundaries (faces) */
    130:       DMStagCreateCompatibleDMStag(ctx.dm_velocity,0,0,1,0,&ctx.dm_buoyancy);
    131:       break;
    132:     default: SETERRQ1(PETSC_COMM_WORLD,PETSC_ERR_SUP,"Not Implemented for dimension %D",ctx.dim);
    133:   }
    134:   DMSetUp(ctx.dm_buoyancy);

    136:   switch (ctx.dim) {
    137:     case 2:
    138:       /* mu and lambda + 2*mu on element centers, mu on corners */
    139:       DMStagCreateCompatibleDMStag(ctx.dm_velocity,1,0,2,0,&ctx.dm_lame);
    140:       break;
    141:     case 3:
    142:       /* mu and lambda + 2*mu on element centers, mu on edges */
    143:       DMStagCreateCompatibleDMStag(ctx.dm_velocity,0,1,0,2,&ctx.dm_lame);
    144:       break;
    145:     default: SETERRQ1(PETSC_COMM_WORLD,PETSC_ERR_SUP,"Not Implemented for dimension %D",ctx.dim);
    146:   }
    147:   DMSetUp(ctx.dm_lame);

    149:   /* Print out some info */
    150:   {
    151:     PetscInt    N[3];
    152:     PetscScalar dx,Vp;

    154:     DMStagGetGlobalSizes(ctx.dm_velocity,&N[0],&N[1],&N[2]);
    155:     dx = (ctx.xmax - ctx.xmin)/N[0];
    156:     Vp = PetscSqrtScalar((ctx.lambda + 2 * ctx.mu) / ctx.rho);
    157:     if (ctx.dim == 2) {
    158:       PetscPrintf(PETSC_COMM_WORLD,"Using a %D x %D mesh\n",N[0],N[1]);
    159:     } else if (ctx.dim == 3) {
    160:       PetscPrintf(PETSC_COMM_WORLD,"Using a %D x %D x %D mesh\n",N[0],N[1],N[2]);
    161:     }
    162:     PetscPrintf(PETSC_COMM_WORLD,"dx: %g\n",dx);
    163:     PetscPrintf(PETSC_COMM_WORLD,"dt: %g\n",ctx.dt);
    164:     PetscPrintf(PETSC_COMM_WORLD,"P-wave velocity: %g\n",PetscSqrtScalar((ctx.lambda + 2 * ctx.mu) / ctx.rho));
    165:     PetscPrintf(PETSC_COMM_WORLD,"V_p dt / dx: %g\n",Vp * ctx.dt / dx);
    166:   }

    168:   /* Populate the coefficient arrays */
    169:   CreateLame(&ctx);
    170:   DMCreateGlobalVector(ctx.dm_buoyancy,&ctx.buoyancy);
    171:   VecSet(ctx.buoyancy,1.0/ctx.rho);

    173:   /* Create vectors to store the system state */
    174:   DMCreateGlobalVector(ctx.dm_velocity,&velocity);
    175:   DMCreateGlobalVector(ctx.dm_stress,&stress);

    177:   /* Initial State */
    178:   VecSet(velocity,0.0);
    179:   VecSet(stress,0.0);
    180:   ForceStress(&ctx,stress,0.0);
    181:   if (ctx.dump_output) {
    182:     DumpVelocity(&ctx,velocity,0);
    183:     DumpStress(&ctx,stress,0);
    184:   }

    186:   /* Time Loop */
    187:   for (timestep = 1; timestep <= ctx.timesteps; ++timestep) {
    188:     const PetscReal t = timestep * ctx.dt;

    190:     UpdateVelocity(&ctx,velocity,stress,ctx.buoyancy);
    191:     UpdateStress(&ctx,velocity,stress,ctx.lame);
    192:     ForceStress(&ctx,stress,t);
    193:     PetscPrintf(PETSC_COMM_WORLD,"Timestep %d, t = %g\n",timestep,(double)t);
    194:     if (ctx.dump_output) {
    195:       DumpVelocity(&ctx,velocity,timestep);
    196:       DumpStress(&ctx,stress,timestep);
    197:     }
    198:   }

    200:   /* Clean up and finalize PETSc */
    201:   VecDestroy(&velocity);
    202:   VecDestroy(&stress);
    203:   VecDestroy(&ctx.lame);
    204:   VecDestroy(&ctx.buoyancy);
    205:   DMDestroy(&ctx.dm_velocity);
    206:   DMDestroy(&ctx.dm_stress);
    207:   DMDestroy(&ctx.dm_buoyancy);
    208:   DMDestroy(&ctx.dm_lame);
    209:   PetscFinalize();
    210:   return ierr;
    211: }

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

    219:   DMCreateGlobalVector(ctx->dm_lame,&ctx->lame);
    220:   DMStagGetGlobalSizes(ctx->dm_lame,&N[0],&N[1],&N[2]);
    221:   DMStagGetCorners(ctx->dm_buoyancy,&startx,&starty,&startz,&nx,&ny,&nz,&extrax,&extray,&extraz);
    222:   if (ctx->dim == 2) {
    223:     /* Element values */
    224:     for (ey=starty; ey<starty+ny; ++ey) {
    225:       for (ex=startx; ex<startx+nx; ++ex) {
    226:         DMStagStencil pos;

    228:         pos.i = ex; pos.j = ey; pos.c = 0; pos.loc = DMSTAG_ELEMENT;
    229:         DMStagVecSetValuesStencil(ctx->dm_lame,ctx->lame,1,&pos,&ctx->lambda,INSERT_VALUES);
    230:         pos.i = ex; pos.j = ey; pos.c = 1; pos.loc = DMSTAG_ELEMENT;
    231:         DMStagVecSetValuesStencil(ctx->dm_lame,ctx->lame,1,&pos,&ctx->mu,INSERT_VALUES);
    232:       }
    233:     }
    234:     /* Vertex Values */
    235:     for (ey=starty; ey<starty+ny+extray; ++ey) {
    236:       for (ex=startx; ex<startx+nx+extrax; ++ex) {
    237:         DMStagStencil pos;

    239:         pos.i = ex; pos.j = ey; pos.c = 0; pos.loc = DMSTAG_DOWN_LEFT;
    240:         DMStagVecSetValuesStencil(ctx->dm_lame,ctx->lame,1,&pos,&ctx->mu,INSERT_VALUES);
    241:       }
    242:     }
    243:   } else if (ctx->dim == 3) {
    244:     /* Element values */
    245:     for (ez=startz; ez<startz+nz; ++ez) {
    246:       for (ey=starty; ey<starty+ny; ++ey) {
    247:         for (ex=startx; ex<startx+nx; ++ex) {
    248:         DMStagStencil pos;

    250:         pos.i = ex; pos.j = ey; pos.k = ez; pos.c = 0; pos.loc = DMSTAG_ELEMENT;
    251:         DMStagVecSetValuesStencil(ctx->dm_lame,ctx->lame,1,&pos,&ctx->lambda,INSERT_VALUES);
    252:         pos.i = ex; pos.j = ey; pos.k = ez; pos.c = 1; pos.loc = DMSTAG_ELEMENT;
    253:         DMStagVecSetValuesStencil(ctx->dm_lame,ctx->lame,1,&pos,&ctx->mu,INSERT_VALUES);
    254:         }
    255:       }
    256:     }
    257:     /* Edge Values */
    258:     for (ez=startz; ez<startz+nz+extraz; ++ez) {
    259:       for (ey=starty; ey<starty+ny+extray; ++ey) {
    260:         for (ex=startx; ex<startx+nx+extrax; ++ex) {
    261:           DMStagStencil   pos;

    263:           if (ex < N[0]) {
    264:             pos.i = ex; pos.j = ey; pos.k = ez; pos.c = 0; pos.loc = DMSTAG_BACK_DOWN;
    265:             DMStagVecSetValuesStencil(ctx->dm_lame,ctx->lame,1,&pos,&ctx->mu,INSERT_VALUES);
    266:           }
    267:           if (ey < N[1]) {
    268:             pos.i = ex; pos.j = ey; pos.k = ez; pos.c = 0; pos.loc = DMSTAG_BACK_LEFT;
    269:             DMStagVecSetValuesStencil(ctx->dm_lame,ctx->lame,1,&pos,&ctx->mu,INSERT_VALUES);
    270:           }
    271:           if (ez < N[2]) {
    272:             pos.i = ex; pos.j = ey; pos.k = ez; pos.c = 0; pos.loc = DMSTAG_DOWN_LEFT;
    273:             DMStagVecSetValuesStencil(ctx->dm_lame,ctx->lame,1,&pos,&ctx->mu,INSERT_VALUES);
    274:           }
    275:         }
    276:       }
    277:     }
    278:   } else SETERRQ1(PetscObjectComm((PetscObject)ctx->dm_velocity),PETSC_ERR_SUP,"Unsupported dim %d",ctx->dim);
    279:   VecAssemblyBegin(ctx->lame);
    280:   VecAssemblyEnd(ctx->lame);
    281:   return(0);
    282: }

    284: static PetscErrorCode ForceStress(const Ctx *ctx,Vec stress, PetscReal t)
    285: {
    286:   PetscErrorCode    ierr;
    287:   PetscInt          start[3],n[3],N[3];
    288:   DMStagStencil     pos;
    289:   PetscBool         this_rank;
    290:   const PetscScalar val = PetscExpReal(-100.0 * t);

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

    296:   /* Normal stresses at a single point */
    297:   this_rank = (PetscBool) (start[0] <= N[0]/2 && N[0]/2 <= start[0] + n[0]);
    298:   this_rank = (PetscBool) (this_rank && start[1] <= N[1]/2 && N[1]/2 <= start[1] + n[1]);
    299:   if (ctx->dim == 3) this_rank = (PetscBool) (this_rank && start[2] <= N[2]/2 && N[2]/2 <= start[2] + n[2]);
    300:   if (this_rank) {
    301:     /* Note integer division to pick element near the center */
    302:     pos.i = N[0]/2; pos.j = N[1]/2; pos.k = N[2]/2; pos.c = 0; pos.loc = DMSTAG_ELEMENT;
    303:     DMStagVecSetValuesStencil(ctx->dm_stress,stress,1,&pos,&val,INSERT_VALUES);
    304:     pos.i = N[0]/2; pos.j = N[1]/2; pos.k = N[2]/2; pos.c = 1; pos.loc = DMSTAG_ELEMENT;
    305:     DMStagVecSetValuesStencil(ctx->dm_stress,stress,1,&pos,&val,INSERT_VALUES);
    306:     if (ctx->dim == 3) {
    307:       pos.i = N[0]/2; pos.j = N[1]/2; pos.k = N[2]/2; pos.c = 2; pos.loc = DMSTAG_ELEMENT;
    308:       DMStagVecSetValuesStencil(ctx->dm_stress,stress,1,&pos,&val,INSERT_VALUES);
    309:     }
    310:   }

    312:   VecAssemblyBegin(stress);
    313:   VecAssemblyEnd(stress);
    314:   return(0);
    315: }

    317: static PetscErrorCode UpdateVelocity_2d(const Ctx *ctx,Vec velocity,Vec stress, Vec buoyancy)
    318: {
    319:   PetscErrorCode    ierr;
    320:   Vec               velocity_local,stress_local,buoyancy_local;
    321:   PetscInt          ex,ey,startx,starty,nx,ny;
    322:   PetscInt          slot_coord_next,slot_coord_element,slot_coord_prev;
    323:   PetscInt          slot_vx_left,slot_vy_down,slot_buoyancy_down,slot_buoyancy_left;
    324:   PetscInt          slot_txx,slot_tyy,slot_txy_downleft,slot_txy_downright,slot_txy_upleft;
    325:   const PetscScalar **arr_coord_x,**arr_coord_y;
    326:   const PetscScalar ***arr_stress,***arr_buoyancy;
    327:   PetscScalar       ***arr_velocity;

    331:   /* Prepare direct access to buoyancy data */
    332:   DMStagGetLocationSlot(ctx->dm_buoyancy,DMSTAG_LEFT,0,&slot_buoyancy_left);
    333:   DMStagGetLocationSlot(ctx->dm_buoyancy,DMSTAG_DOWN,0,&slot_buoyancy_down);
    334:   DMGetLocalVector(ctx->dm_buoyancy,&buoyancy_local);
    335:   DMGlobalToLocal(ctx->dm_buoyancy,buoyancy,INSERT_VALUES,buoyancy_local);
    336:   DMStagVecGetArrayRead(ctx->dm_buoyancy,buoyancy_local,&arr_buoyancy);

    338:   /* Prepare read-only access to stress data */
    339:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_ELEMENT,   0,&slot_txx);
    340:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_ELEMENT,   1,&slot_tyy);
    341:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_UP_LEFT,   0,&slot_txy_upleft);
    342:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_DOWN_LEFT, 0,&slot_txy_downleft);
    343:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_DOWN_RIGHT,0,&slot_txy_downright);
    344:   DMGetLocalVector(ctx->dm_stress,&stress_local);
    345:   DMGlobalToLocal(ctx->dm_stress,stress,INSERT_VALUES,stress_local);
    346:   DMStagVecGetArrayRead(ctx->dm_stress,stress_local,&arr_stress);

    348:   /* Prepare read-write access to velocity data */
    349:   DMStagGetLocationSlot(ctx->dm_velocity,DMSTAG_LEFT,0,&slot_vx_left);
    350:   DMStagGetLocationSlot(ctx->dm_velocity,DMSTAG_DOWN,0,&slot_vy_down);
    351:   DMGetLocalVector(ctx->dm_velocity,&velocity_local);
    352:   DMGlobalToLocal(ctx->dm_velocity,velocity,INSERT_VALUES,velocity_local);
    353:   DMStagVecGetArray(ctx->dm_velocity,velocity_local,&arr_velocity);

    355:   /* Prepare read-only access to coordinate data */
    356:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity,DMSTAG_LEFT,&slot_coord_prev);
    357:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity,DMSTAG_RIGHT,&slot_coord_next);
    358:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity,DMSTAG_ELEMENT,&slot_coord_element);
    359:   DMStagGetProductCoordinateArrays(ctx->dm_velocity,&arr_coord_x,&arr_coord_y,NULL);

    361:   /* Iterate over interior of the domain, updating the velocities */
    362:   DMStagGetCorners(ctx->dm_velocity,&startx,&starty,NULL,&nx,&ny,NULL,NULL,NULL,NULL);
    363:   for (ey = starty; ey < starty + ny; ++ey) {
    364:     for (ex = startx; ex < startx + nx; ++ex) {

    366:       /* Update y-velocity */
    367:       if (ey > 0) {
    368:         const PetscScalar dx = arr_coord_x[ex][slot_coord_next]    - arr_coord_x[ex  ][slot_coord_prev];
    369:         const PetscScalar dy = arr_coord_y[ey][slot_coord_element] - arr_coord_y[ey-1][slot_coord_element];
    370:         const PetscScalar B = arr_buoyancy[ey][ex][slot_buoyancy_down];

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

    377:       /* Update x-velocity */
    378:       if (ex > 0) {
    379:         const PetscScalar dx = arr_coord_x[ex][slot_coord_element] - arr_coord_x[ex-1][slot_coord_element];
    380:         const PetscScalar dy = arr_coord_y[ey][slot_coord_next]    - arr_coord_y[ey  ][slot_coord_prev];
    381:         const PetscScalar B = arr_buoyancy[ey][ex][slot_buoyancy_left];

    383:         arr_velocity[ey][ex][slot_vx_left] += B * ctx->dt * (
    384:               (arr_stress[ey][ex][slot_txx]        - arr_stress[ey][ex-1][slot_txx]) / dx
    385:             + (arr_stress[ey][ex][slot_txy_upleft] - arr_stress[ey][ex  ][slot_txy_downleft]) / dy);
    386:       }
    387:     }
    388:   }

    390:   /* Restore all access */
    391:   DMStagVecRestoreArrayRead(ctx->dm_buoyancy,buoyancy_local,&arr_buoyancy);
    392:   DMRestoreLocalVector(ctx->dm_buoyancy,&buoyancy_local);
    393:   DMStagVecRestoreArrayRead(ctx->dm_stress,stress_local,&arr_stress);
    394:   DMRestoreLocalVector(ctx->dm_stress,&stress_local);
    395:   DMStagVecRestoreArray(ctx->dm_velocity,velocity_local,&arr_velocity);
    396:   DMLocalToGlobal(ctx->dm_velocity,velocity_local,INSERT_VALUES,velocity);
    397:   DMRestoreLocalVector(ctx->dm_velocity,&velocity_local);
    398:   DMStagRestoreProductCoordinateArrays(ctx->dm_velocity,&arr_coord_x,&arr_coord_y,NULL);
    399:   return(0);
    400: }

    402: static PetscErrorCode UpdateVelocity_3d(const Ctx *ctx,Vec velocity,Vec stress, Vec buoyancy)
    403: {
    404:   PetscErrorCode    ierr;
    405:   Vec               velocity_local,stress_local,buoyancy_local;
    406:   PetscInt          ex,ey,ez,startx,starty,startz,nx,ny,nz;
    407:   PetscInt          slot_coord_next,slot_coord_element,slot_coord_prev;
    408:   PetscInt          slot_vx_left,slot_vy_down,slot_vz_back,slot_buoyancy_down,slot_buoyancy_left,slot_buoyancy_back;
    409:   PetscInt          slot_txx,slot_tyy,slot_tzz;
    410:   PetscInt          slot_txy_downleft,slot_txy_downright,slot_txy_upleft;
    411:   PetscInt          slot_txz_backleft,slot_txz_backright,slot_txz_frontleft;
    412:   PetscInt          slot_tyz_backdown,slot_tyz_frontdown,slot_tyz_backup;
    413:   const PetscScalar **arr_coord_x,**arr_coord_y,**arr_coord_z;
    414:   const PetscScalar ****arr_stress,****arr_buoyancy;
    415:   PetscScalar       ****arr_velocity;

    419:   /* Prepare direct access to buoyancy data */
    420:   DMStagGetLocationSlot(ctx->dm_buoyancy,DMSTAG_LEFT,0,&slot_buoyancy_left);
    421:   DMStagGetLocationSlot(ctx->dm_buoyancy,DMSTAG_DOWN,0,&slot_buoyancy_down);
    422:   DMStagGetLocationSlot(ctx->dm_buoyancy,DMSTAG_BACK,0,&slot_buoyancy_back);
    423:   DMGetLocalVector(ctx->dm_buoyancy,&buoyancy_local);
    424:   DMGlobalToLocal(ctx->dm_buoyancy,buoyancy,INSERT_VALUES,buoyancy_local);
    425:   DMStagVecGetArrayRead(ctx->dm_buoyancy,buoyancy_local,&arr_buoyancy);

    427:   /* Prepare read-only access to stress data */
    428:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_ELEMENT,   0,&slot_txx);
    429:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_ELEMENT,   1,&slot_tyy);
    430:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_ELEMENT,   2,&slot_tzz);
    431:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_UP_LEFT,   0,&slot_txy_upleft);
    432:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_DOWN_LEFT, 0,&slot_txy_downleft);
    433:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_DOWN_RIGHT,0,&slot_txy_downright);
    434:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_BACK_LEFT, 0,&slot_txz_backleft);
    435:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_BACK_RIGHT,0,&slot_txz_backright);
    436:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_FRONT_LEFT,0,&slot_txz_frontleft);
    437:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_BACK_DOWN, 0,&slot_tyz_backdown);
    438:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_BACK_UP,   0,&slot_tyz_backup);
    439:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_FRONT_DOWN,0,&slot_tyz_frontdown);
    440:   DMGetLocalVector(ctx->dm_stress,&stress_local);
    441:   DMGlobalToLocal(ctx->dm_stress,stress,INSERT_VALUES,stress_local);
    442:   DMStagVecGetArrayRead(ctx->dm_stress,stress_local,&arr_stress);

    444:   /* Prepare read-write access to velocity data */
    445:   DMStagGetLocationSlot(ctx->dm_velocity,DMSTAG_LEFT,0,&slot_vx_left);
    446:   DMStagGetLocationSlot(ctx->dm_velocity,DMSTAG_DOWN,0,&slot_vy_down);
    447:   DMStagGetLocationSlot(ctx->dm_velocity,DMSTAG_BACK,0,&slot_vz_back);
    448:   DMGetLocalVector(ctx->dm_velocity,&velocity_local);
    449:   DMGlobalToLocal(ctx->dm_velocity,velocity,INSERT_VALUES,velocity_local);
    450:   DMStagVecGetArray(ctx->dm_velocity,velocity_local,&arr_velocity);

    452:   /* Prepare read-only access to coordinate data */
    453:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity,DMSTAG_LEFT,&slot_coord_prev);
    454:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity,DMSTAG_RIGHT,&slot_coord_next);
    455:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity,DMSTAG_ELEMENT,&slot_coord_element);
    456:   DMStagGetProductCoordinateArrays(ctx->dm_velocity,&arr_coord_x,&arr_coord_y,&arr_coord_z);

    458:   /* Iterate over interior of the domain, updating the velocities */
    459:   DMStagGetCorners(ctx->dm_velocity,&startx,&starty,&startz,&nx,&ny,&nz,NULL,NULL,NULL);
    460:   for (ez = startz; ez < startz + nz; ++ez) {
    461:     for (ey = starty; ey < starty + ny; ++ey) {
    462:       for (ex = startx; ex < startx + nx; ++ex) {

    464:         /* Update y-velocity */
    465:         if (ey > 0) {
    466:           const PetscScalar dx = arr_coord_x[ex][slot_coord_next]    - arr_coord_x[ex  ][slot_coord_prev];
    467:           const PetscScalar dy = arr_coord_y[ey][slot_coord_element] - arr_coord_y[ey-1][slot_coord_element];
    468:           const PetscScalar dz = arr_coord_z[ez][slot_coord_next]    - arr_coord_z[ez  ][slot_coord_prev];
    469:           const PetscScalar B = arr_buoyancy[ez][ey][ex][slot_buoyancy_down];

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

    477:         /* Update x-velocity */
    478:         if (ex > 0) {
    479:           const PetscScalar dx = arr_coord_x[ex][slot_coord_element] - arr_coord_x[ex-1][slot_coord_element];
    480:           const PetscScalar dy = arr_coord_y[ey][slot_coord_next]    - arr_coord_y[ey  ][slot_coord_prev];
    481:           const PetscScalar dz = arr_coord_z[ez][slot_coord_next]    - arr_coord_z[ez  ][slot_coord_prev];
    482:           const PetscScalar B = arr_buoyancy[ez][ey][ex][slot_buoyancy_left];

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

    490:         /* Update z-velocity */
    491:         if (ez > 0) {
    492:           const PetscScalar dx = arr_coord_x[ex][slot_coord_next]    - arr_coord_x[ex  ][slot_coord_prev];
    493:           const PetscScalar dy = arr_coord_y[ey][slot_coord_next]    - arr_coord_y[ey  ][slot_coord_prev];
    494:           const PetscScalar dz = arr_coord_z[ez][slot_coord_element] - arr_coord_z[ez-1][slot_coord_element];
    495:           const PetscScalar B = arr_buoyancy[ez][ey][ex][slot_buoyancy_back];

    497:           arr_velocity[ez][ey][ex][slot_vz_back] += B * ctx->dt * (
    498:               (arr_stress[ez][ey][ex][slot_txz_backright] - arr_stress[ez  ][ey][ex][slot_txz_backleft]) / dx
    499:             + (arr_stress[ez][ey][ex][slot_tyz_backup]    - arr_stress[ez  ][ey][ex][slot_tyz_backdown]) / dy
    500:             + (arr_stress[ez][ey][ex][slot_tzz]           - arr_stress[ez-1][ey][ex][slot_tzz]) / dz);
    501:         }
    502:       }
    503:     }
    504:   }

    506:   /* Restore all access */
    507:   DMStagVecRestoreArrayRead(ctx->dm_buoyancy,buoyancy_local,&arr_buoyancy);
    508:   DMRestoreLocalVector(ctx->dm_buoyancy,&buoyancy_local);
    509:   DMStagVecRestoreArrayRead(ctx->dm_stress,stress_local,&arr_stress);
    510:   DMRestoreLocalVector(ctx->dm_stress,&stress_local);
    511:   DMStagVecRestoreArray(ctx->dm_velocity,velocity_local,&arr_velocity);
    512:   DMLocalToGlobal(ctx->dm_velocity,velocity_local,INSERT_VALUES,velocity);
    513:   DMRestoreLocalVector(ctx->dm_velocity,&velocity_local);
    514:   DMStagRestoreProductCoordinateArrays(ctx->dm_velocity,&arr_coord_x,&arr_coord_y,&arr_coord_z);
    515:   return(0);
    516: }

    518: static PetscErrorCode UpdateVelocity(const Ctx *ctx,Vec velocity,Vec stress, Vec buoyancy)
    519: {

    523:   if (ctx->dim == 2) {
    524:     UpdateVelocity_2d(ctx,velocity,stress,buoyancy);
    525:   } else if (ctx->dim == 3) {
    526:     UpdateVelocity_3d(ctx,velocity,stress,buoyancy);
    527:   } else SETERRQ1(PetscObjectComm((PetscObject)ctx->dm_velocity),PETSC_ERR_SUP,"Unsupported dim %d",ctx->dim);
    528:   return(0);
    529: }

    531: static PetscErrorCode UpdateStress_2d(const Ctx *ctx,Vec velocity,Vec stress, Vec lame)
    532: {
    533:   PetscErrorCode    ierr;
    534:   Vec               velocity_local,stress_local,lame_local;
    535:   PetscInt          ex,ey,startx,starty,nx,ny;
    536:   PetscInt          slot_coord_next,slot_coord_element,slot_coord_prev;
    537:   PetscInt          slot_vx_left,slot_vy_down,slot_vx_right,slot_vy_up;
    538:   PetscInt          slot_mu_element,slot_lambda_element,slot_mu_downleft;
    539:   PetscInt          slot_txx,slot_tyy,slot_txy_downleft;
    540:   const PetscScalar **arr_coord_x,**arr_coord_y;
    541:   const PetscScalar ***arr_velocity,***arr_lame;
    542:   PetscScalar       ***arr_stress;

    546:   /* Prepare read-write access to stress data */
    547:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_ELEMENT,0,&slot_txx);
    548:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_ELEMENT,1,&slot_tyy);
    549:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_DOWN_LEFT,0,&slot_txy_downleft);
    550:   DMGetLocalVector(ctx->dm_stress,&stress_local);
    551:   DMGlobalToLocal(ctx->dm_stress,stress,INSERT_VALUES,stress_local);
    552:   DMStagVecGetArray(ctx->dm_stress,stress_local,&arr_stress);

    554:   /* Prepare read-only access to velocity data */
    555:   DMStagGetLocationSlot(ctx->dm_velocity,DMSTAG_LEFT,0,&slot_vx_left);
    556:   DMStagGetLocationSlot(ctx->dm_velocity,DMSTAG_RIGHT,0,&slot_vx_right);
    557:   DMStagGetLocationSlot(ctx->dm_velocity,DMSTAG_DOWN,0,&slot_vy_down);
    558:   DMStagGetLocationSlot(ctx->dm_velocity,DMSTAG_UP,0,&slot_vy_up);
    559:   DMGetLocalVector(ctx->dm_velocity,&velocity_local);
    560:   DMGlobalToLocal(ctx->dm_velocity,velocity,INSERT_VALUES,velocity_local);
    561:   DMStagVecGetArrayRead(ctx->dm_velocity,velocity_local,&arr_velocity);

    563:   /* Prepare read-only access to Lame' data */
    564:   DMStagGetLocationSlot(ctx->dm_lame,DMSTAG_ELEMENT,0,&slot_lambda_element);
    565:   DMStagGetLocationSlot(ctx->dm_lame,DMSTAG_ELEMENT,1,&slot_mu_element);
    566:   DMStagGetLocationSlot(ctx->dm_lame,DMSTAG_DOWN_LEFT,0,&slot_mu_downleft);
    567:   DMGetLocalVector(ctx->dm_lame,&lame_local);
    568:   DMGlobalToLocal(ctx->dm_lame,lame,INSERT_VALUES,lame_local);
    569:   DMStagVecGetArrayRead(ctx->dm_lame,lame_local,&arr_lame);

    571:   /* Prepare read-only access to coordinate data */
    572:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity,DMSTAG_LEFT,&slot_coord_prev);
    573:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity,DMSTAG_RIGHT,&slot_coord_next);
    574:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity,DMSTAG_ELEMENT,&slot_coord_element);
    575:   DMStagGetProductCoordinateArrays(ctx->dm_velocity,&arr_coord_x,&arr_coord_y,NULL);

    577:   /* Iterate over the interior of the domain, updating the stresses */
    578:   DMStagGetCorners(ctx->dm_velocity,&startx,&starty,NULL,&nx,&ny,NULL,NULL,NULL,NULL);
    579:   for (ey = starty; ey < starty + ny; ++ey) {
    580:     for (ex = startx; ex < startx + nx; ++ex) {

    582:       /* Update tau_xx and tau_yy*/
    583:       {
    584:         const PetscScalar dx   = arr_coord_x[ex][slot_coord_next] - arr_coord_x[ex][slot_coord_prev];
    585:         const PetscScalar dy   = arr_coord_y[ey][slot_coord_next] - arr_coord_y[ey][slot_coord_prev];
    586:         const PetscScalar L    = arr_lame[ey][ex][slot_lambda_element];
    587:         const PetscScalar Lp2M = arr_lame[ey][ex][slot_lambda_element] + 2 * arr_lame[ey][ex][slot_mu_element];

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

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

    598:       /* Update tau_xy */
    599:       if (ex > 0 && ey > 0) {
    600:         const PetscScalar dx = arr_coord_x[ex][slot_coord_element] - arr_coord_x[ex-1][slot_coord_element];
    601:         const PetscScalar dy = arr_coord_y[ey][slot_coord_element] - arr_coord_y[ey-1][slot_coord_element];
    602:         const PetscScalar M  = arr_lame[ey][ex][slot_mu_downleft];

    604:         arr_stress[ey][ex][slot_txy_downleft] += M * ctx->dt * (
    605:             (arr_velocity[ey][ex][slot_vx_left] - arr_velocity[ey-1][ex][slot_vx_left]) / dy
    606:           + (arr_velocity[ey][ex][slot_vy_down] - arr_velocity[ey][ex-1][slot_vy_down]) / dx);
    607:       }
    608:     }
    609:   }

    611:   /* Restore all access */
    612:   DMStagVecRestoreArray(ctx->dm_stress,stress_local,&arr_stress);
    613:   DMLocalToGlobal(ctx->dm_stress,stress_local,INSERT_VALUES,stress);
    614:   DMRestoreLocalVector(ctx->dm_stress,&stress_local);
    615:   DMStagVecRestoreArrayRead(ctx->dm_velocity,velocity_local,&arr_velocity);
    616:   DMRestoreLocalVector(ctx->dm_velocity,&velocity_local);
    617:   DMStagVecRestoreArrayRead(ctx->dm_lame,lame_local,&arr_lame);
    618:   DMRestoreLocalVector(ctx->dm_lame,&lame_local);
    619:   DMStagRestoreProductCoordinateArrays(ctx->dm_velocity,&arr_coord_x,&arr_coord_y,NULL);
    620:   return(0);
    621: }

    623: static PetscErrorCode UpdateStress_3d(const Ctx *ctx,Vec velocity,Vec stress, Vec lame)
    624: {
    625:   PetscErrorCode    ierr;
    626:   Vec               velocity_local,stress_local,lame_local;
    627:   PetscInt          ex,ey,ez,startx,starty,startz,nx,ny,nz;
    628:   PetscInt          slot_coord_next,slot_coord_element,slot_coord_prev;
    629:   PetscInt          slot_vx_left,slot_vy_down,slot_vx_right,slot_vy_up,slot_vz_back,slot_vz_front;
    630:   PetscInt          slot_mu_element,slot_lambda_element,slot_mu_downleft,slot_mu_backdown,slot_mu_backleft;
    631:   PetscInt          slot_txx,slot_tyy,slot_tzz,slot_txy_downleft,slot_txz_backleft,slot_tyz_backdown;
    632:   const PetscScalar **arr_coord_x,**arr_coord_y,**arr_coord_z;
    633:   const PetscScalar ****arr_velocity,****arr_lame;
    634:   PetscScalar       ****arr_stress;

    638:   /* Prepare read-write access to stress data */
    639:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_ELEMENT,0,&slot_txx);
    640:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_ELEMENT,1,&slot_tyy);
    641:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_ELEMENT,2,&slot_tzz);
    642:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_DOWN_LEFT,0,&slot_txy_downleft);
    643:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_BACK_LEFT,0,&slot_txz_backleft);
    644:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_BACK_DOWN,0,&slot_tyz_backdown);
    645:   DMGetLocalVector(ctx->dm_stress,&stress_local);
    646:   DMGlobalToLocal(ctx->dm_stress,stress,INSERT_VALUES,stress_local);
    647:   DMStagVecGetArray(ctx->dm_stress,stress_local,&arr_stress);

    649:   /* Prepare read-only access to velocity data */
    650:   DMStagGetLocationSlot(ctx->dm_velocity,DMSTAG_LEFT, 0,&slot_vx_left);
    651:   DMStagGetLocationSlot(ctx->dm_velocity,DMSTAG_RIGHT,0,&slot_vx_right);
    652:   DMStagGetLocationSlot(ctx->dm_velocity,DMSTAG_DOWN, 0,&slot_vy_down);
    653:   DMStagGetLocationSlot(ctx->dm_velocity,DMSTAG_UP,   0,&slot_vy_up);
    654:   DMStagGetLocationSlot(ctx->dm_velocity,DMSTAG_BACK, 0,&slot_vz_back);
    655:   DMStagGetLocationSlot(ctx->dm_velocity,DMSTAG_FRONT,0,&slot_vz_front);
    656:   DMGetLocalVector(ctx->dm_velocity,&velocity_local);
    657:   DMGlobalToLocal(ctx->dm_velocity,velocity,INSERT_VALUES,velocity_local);
    658:   DMStagVecGetArrayRead(ctx->dm_velocity,velocity_local,&arr_velocity);

    660:   /* Prepare read-only access to Lame' data */
    661:   DMStagGetLocationSlot(ctx->dm_lame,DMSTAG_ELEMENT,  0,&slot_lambda_element);
    662:   DMStagGetLocationSlot(ctx->dm_lame,DMSTAG_ELEMENT,  1,&slot_mu_element);
    663:   DMStagGetLocationSlot(ctx->dm_lame,DMSTAG_DOWN_LEFT,0,&slot_mu_downleft);
    664:   DMStagGetLocationSlot(ctx->dm_lame,DMSTAG_BACK_LEFT,0,&slot_mu_backleft);
    665:   DMStagGetLocationSlot(ctx->dm_lame,DMSTAG_BACK_DOWN,0,&slot_mu_backdown);
    666:   DMGetLocalVector(ctx->dm_lame,&lame_local);
    667:   DMGlobalToLocal(ctx->dm_lame,lame,INSERT_VALUES,lame_local);
    668:   DMStagVecGetArrayRead(ctx->dm_lame,lame_local,&arr_lame);

    670:   /* Prepare read-only access to coordinate data */
    671:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity,DMSTAG_LEFT,&slot_coord_prev);
    672:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity,DMSTAG_RIGHT,&slot_coord_next);
    673:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity,DMSTAG_ELEMENT,&slot_coord_element);
    674:   DMStagGetProductCoordinateArrays(ctx->dm_velocity,&arr_coord_x,&arr_coord_y,&arr_coord_z);

    676:   /* Iterate over the interior of the domain, updating the stresses */
    677:   DMStagGetCorners(ctx->dm_velocity,&startx,&starty,&startz,&nx,&ny,&nz,NULL,NULL,NULL);
    678:   for (ez = startz; ez < startz + nz; ++ez) {
    679:     for (ey = starty; ey < starty + ny; ++ey) {
    680:       for (ex = startx; ex < startx + nx; ++ex) {

    682:         /* Update tau_xx, tau_yy, and tau_zz*/
    683:         {
    684:           const PetscScalar dx   = arr_coord_x[ex][slot_coord_next] - arr_coord_x[ex][slot_coord_prev];
    685:           const PetscScalar dy   = arr_coord_y[ey][slot_coord_next] - arr_coord_y[ey][slot_coord_prev];
    686:           const PetscScalar dz   = arr_coord_z[ez][slot_coord_next] - arr_coord_z[ez][slot_coord_prev];
    687:           const PetscScalar L    = arr_lame[ez][ey][ex][slot_lambda_element];
    688:           const PetscScalar Lp2M = arr_lame[ez][ey][ex][slot_lambda_element] + 2 * arr_lame[ez][ey][ex][slot_mu_element];

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

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

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

    706:         /* Update tau_xy, tau_xz, tau_yz */
    707:         {
    708:           PetscScalar dx,dy,dz;

    710:           if (ex > 0)  dx = arr_coord_x[ex][slot_coord_element] - arr_coord_x[ex-1][slot_coord_element];
    711:           if (ey > 0)  dy = arr_coord_y[ey][slot_coord_element] - arr_coord_y[ey-1][slot_coord_element];
    712:           if (ez > 0)  dz = arr_coord_z[ez][slot_coord_element] - arr_coord_z[ez-1][slot_coord_element];

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

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

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

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

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

    735:             arr_stress[ez][ey][ex][slot_tyz_backdown] += M * ctx->dt * (
    736:                 (arr_velocity[ez][ey][ex][slot_vy_down] - arr_velocity[ez-1][ey][ex][slot_vy_down]) / dz
    737:               + (arr_velocity[ez][ey][ex][slot_vz_back] - arr_velocity[ez][ey-1][ex][slot_vz_back]) / dy);
    738:           }
    739:         }
    740:       }
    741:     }
    742:   }

    744:   /* Restore all access */
    745:   DMStagVecRestoreArray(ctx->dm_stress,stress_local,&arr_stress);
    746:   DMLocalToGlobal(ctx->dm_stress,stress_local,INSERT_VALUES,stress);
    747:   DMRestoreLocalVector(ctx->dm_stress,&stress_local);
    748:   DMStagVecRestoreArrayRead(ctx->dm_velocity,velocity_local,&arr_velocity);
    749:   DMRestoreLocalVector(ctx->dm_velocity,&velocity_local);
    750:   DMStagVecRestoreArrayRead(ctx->dm_lame,lame_local,&arr_lame);
    751:   DMRestoreLocalVector(ctx->dm_lame,&lame_local);
    752:   DMStagRestoreProductCoordinateArrays(ctx->dm_velocity,&arr_coord_x,&arr_coord_y,&arr_coord_z);
    753:   return(0);
    754: }

    756: static PetscErrorCode UpdateStress(const Ctx *ctx,Vec velocity,Vec stress, Vec lame)
    757: {

    761:   if (ctx->dim == 2) {
    762:     UpdateStress_2d(ctx,velocity,stress,lame);
    763:   } else if (ctx->dim ==3) {
    764:     UpdateStress_3d(ctx,velocity,stress,lame);
    765:   }
    766:   return(0);
    767: }

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

    777:     DMStagVecSplitToDMDA(ctx->dm_stress,stress,DMSTAG_ELEMENT,-ctx->dim,&da_normal,&vec_normal);
    778:     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:       PetscSNPrintf(filename,sizeof(filename),"ex6_stress_normal_%.4D.vtr",timestep);
    786:       PetscViewerVTKOpen(PetscObjectComm((PetscObject)da_normal),filename,FILE_MODE_WRITE,&viewer);
    787:       VecView(vec_normal,viewer);
    788:       PetscViewerDestroy(&viewer);
    789:       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:     DMStagVecSplitToDMDA(ctx->dm_stress,stress,DMSTAG_DOWN_LEFT,0,&da_shear,&vec_shear);
    795:     PetscObjectSetName((PetscObject)vec_shear,"shear stresses");

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

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

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

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

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

    846:         from[0].i = ex; from[0].j = ey; from[0].loc = DMSTAG_UP;    from[0].c = 0;
    847:         from[1].i = ex; from[1].j = ey; from[1].loc = DMSTAG_DOWN;  from[1].c = 0;
    848:         from[2].i = ex; from[2].j = ey; from[2].loc = DMSTAG_LEFT;  from[2].c = 0;
    849:         from[3].i = ex; from[3].j = ey; from[3].loc = DMSTAG_RIGHT; from[3].c = 0;
    850:         DMStagVecGetValuesStencil(ctx->dm_velocity,velocity_local,4,from,valFrom);
    851:         to[0].i = ex; to[0].j = ey; to[0].loc = DMSTAG_ELEMENT;    to[0].c = 0; valTo[0] = 0.5 * (valFrom[2] + valFrom[3]);
    852:         to[1].i = ex; to[1].j = ey; to[1].loc = DMSTAG_ELEMENT;    to[1].c = 1; valTo[1] = 0.5 * (valFrom[0] + valFrom[1]);
    853:         DMStagVecSetValuesStencil(dmVelAvg,velAvg,2,to,valTo,INSERT_VALUES);
    854:       }
    855:     }
    856:   } else if (ctx->dim == 3) {
    857:     for (ez = startz; ez<startz+nz; ++ez) {
    858:       for (ey = starty; ey<starty+ny; ++ey) {
    859:         for (ex = startx; ex<startx+nx; ++ex) {
    860:           DMStagStencil from[6],to[3];
    861:           PetscScalar   valFrom[6],valTo[3];

    863:           from[0].i = ex; from[0].j = ey; from[0].k = ez; from[0].loc = DMSTAG_UP;    from[0].c = 0;
    864:           from[1].i = ex; from[1].j = ey; from[1].k = ez; from[1].loc = DMSTAG_DOWN;  from[1].c = 0;
    865:           from[2].i = ex; from[2].j = ey; from[2].k = ez; from[2].loc = DMSTAG_LEFT;  from[2].c = 0;
    866:           from[3].i = ex; from[3].j = ey; from[3].k = ez; from[3].loc = DMSTAG_RIGHT; from[3].c = 0;
    867:           from[4].i = ex; from[4].j = ey; from[4].k = ez; from[4].loc = DMSTAG_FRONT; from[4].c = 0;
    868:           from[5].i = ex; from[5].j = ey; from[5].k = ez; from[5].loc = DMSTAG_BACK;  from[5].c = 0;
    869:           DMStagVecGetValuesStencil(ctx->dm_velocity,velocity_local,6,from,valFrom);
    870:           to[0].i = ex; to[0].j = ey; to[0].k = ez; to[0].loc = DMSTAG_ELEMENT;    to[0].c = 0; valTo[0] = 0.5 * (valFrom[2] + valFrom[3]);
    871:           to[1].i = ex; to[1].j = ey; to[1].k = ez; to[1].loc = DMSTAG_ELEMENT;    to[1].c = 1; valTo[1] = 0.5 * (valFrom[0] + valFrom[1]);
    872:           to[2].i = ex; to[2].j = ey; to[2].k = ez; to[2].loc = DMSTAG_ELEMENT;    to[2].c = 2; valTo[2] = 0.5 * (valFrom[4] + valFrom[5]);
    873:           DMStagVecSetValuesStencil(dmVelAvg,velAvg,3,to,valTo,INSERT_VALUES);
    874:         }
    875:       }
    876:     }
    877:   } else SETERRQ1(PetscObjectComm((PetscObject)ctx->dm_velocity),PETSC_ERR_SUP,"Unsupported dim %d",ctx->dim);
    878:   VecAssemblyBegin(velAvg);
    879:   VecAssemblyEnd(velAvg);
    880:   DMRestoreLocalVector(ctx->dm_velocity,&velocity_local);

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

    885:   /* Dump element-based fields to a .vtr file */
    886:   {
    887:     PetscViewer viewer;
    888:     char        filename[PETSC_MAX_PATH_LEN];

    890:     PetscSNPrintf(filename,sizeof(filename),"ex6_velavg_%.4D.vtr",timestep);
    891:     PetscViewerVTKOpen(PetscObjectComm((PetscObject)daVelAvg),filename,FILE_MODE_WRITE,&viewer);
    892:     VecView(vecVelAvg,viewer);
    893:     PetscViewerDestroy(&viewer);
    894:     PetscPrintf(PETSC_COMM_WORLD,"Created %s\n",filename);
    895:   }

    897:   /* Destroy DMDAs and Vecs */
    898:   VecDestroy(&vecVelAvg);
    899:   DMDestroy(&daVelAvg);
    900:   VecDestroy(&velAvg);
    901:   DMDestroy(&dmVelAvg);
    902:   return(0);
    903: }

    905: /*TEST

    907:    test:
    908:       suffix: 1
    909:       requires: !complex
    910:       nsize: 1
    911:       args: -stag_grid_x 12 -stag_grid_y 7 -nsteps 4 -dump_output 0

    913:    test:
    914:       suffix: 2
    915:       requires: !complex
    916:       nsize: 9
    917:       args: -stag_grid_x 12 -stag_grid_y 15 -nsteps 3 -dump_output 0

    919:    test:
    920:       suffix: 3
    921:       requires: !complex
    922:       nsize: 1
    923:       args: -stag_grid_x 12 -stag_grid_y 7 -stag_grid_z 11 -nsteps 2 -dim 3 -dump_output 0

    925:    test:
    926:       suffix: 4
    927:       requires: !complex
    928:       nsize: 12
    929:       args: -stag_grid_x 12 -stag_grid_y 15 -stag_grid_z 8 -nsteps 3 -dim 3 -dump_output 0

    931: TEST*/