Date: 2022-08-02T03:14:22+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: {
     60:   Ctx            ctx;
     61:   Vec            velocity,stress;
     62:   PetscInt       timestep;

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    309:   VecAssemblyBegin(stress);
    310:   VecAssemblyEnd(stress);
    311:   return 0;
    312: }

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

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

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

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

    351:   /* Prepare read-only access to coordinate data */
    352:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity,DMSTAG_LEFT,&slot_coord_prev);
    353:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity,DMSTAG_RIGHT,&slot_coord_next);
    354:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity,DMSTAG_ELEMENT,&slot_coord_element);
    355:   DMStagGetProductCoordinateArrays(ctx->dm_velocity,&arr_coord_x,&arr_coord_y,NULL);

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

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

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

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

    379:         arr_velocity[ey][ex][slot_vx_left] += B * ctx->dt * (
    380:               (arr_stress[ey][ex][slot_txx]        - arr_stress[ey][ex-1][slot_txx]) / dx
    381:             + (arr_stress[ey][ex][slot_txy_upleft] - arr_stress[ey][ex  ][slot_txy_downleft]) / dy);
    382:       }
    383:     }
    384:   }

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

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

    414:   /* Prepare direct access to buoyancy data */
    415:   DMStagGetLocationSlot(ctx->dm_buoyancy,DMSTAG_LEFT,0,&slot_buoyancy_left);
    416:   DMStagGetLocationSlot(ctx->dm_buoyancy,DMSTAG_DOWN,0,&slot_buoyancy_down);
    417:   DMStagGetLocationSlot(ctx->dm_buoyancy,DMSTAG_BACK,0,&slot_buoyancy_back);
    418:   DMGetLocalVector(ctx->dm_buoyancy,&buoyancy_local);
    419:   DMGlobalToLocal(ctx->dm_buoyancy,buoyancy,INSERT_VALUES,buoyancy_local);
    420:   DMStagVecGetArrayRead(ctx->dm_buoyancy,buoyancy_local,&arr_buoyancy);

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

    439:   /* Prepare read-write access to velocity data */
    440:   DMStagGetLocationSlot(ctx->dm_velocity,DMSTAG_LEFT,0,&slot_vx_left);
    441:   DMStagGetLocationSlot(ctx->dm_velocity,DMSTAG_DOWN,0,&slot_vy_down);
    442:   DMStagGetLocationSlot(ctx->dm_velocity,DMSTAG_BACK,0,&slot_vz_back);
    443:   DMGetLocalVector(ctx->dm_velocity,&velocity_local);
    444:   DMGlobalToLocal(ctx->dm_velocity,velocity,INSERT_VALUES,velocity_local);
    445:   DMStagVecGetArray(ctx->dm_velocity,velocity_local,&arr_velocity);

    447:   /* Prepare read-only access to coordinate data */
    448:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity,DMSTAG_LEFT,&slot_coord_prev);
    449:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity,DMSTAG_RIGHT,&slot_coord_next);
    450:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity,DMSTAG_ELEMENT,&slot_coord_element);
    451:   DMStagGetProductCoordinateArrays(ctx->dm_velocity,&arr_coord_x,&arr_coord_y,&arr_coord_z);

    453:   /* Iterate over interior of the domain, updating the velocities */
    454:   DMStagGetCorners(ctx->dm_velocity,&startx,&starty,&startz,&nx,&ny,&nz,NULL,NULL,NULL);
    455:   for (ez = startz; ez < startz + nz; ++ez) {
    456:     for (ey = starty; ey < starty + ny; ++ey) {
    457:       for (ex = startx; ex < startx + nx; ++ex) {

    459:         /* Update y-velocity */
    460:         if (ey > 0) {
    461:           const PetscScalar dx = arr_coord_x[ex][slot_coord_next]    - arr_coord_x[ex  ][slot_coord_prev];
    462:           const PetscScalar dy = arr_coord_y[ey][slot_coord_element] - arr_coord_y[ey-1][slot_coord_element];
    463:           const PetscScalar dz = arr_coord_z[ez][slot_coord_next]    - arr_coord_z[ez  ][slot_coord_prev];
    464:           const PetscScalar B = arr_buoyancy[ez][ey][ex][slot_buoyancy_down];

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

    472:         /* Update x-velocity */
    473:         if (ex > 0) {
    474:           const PetscScalar dx = arr_coord_x[ex][slot_coord_element] - arr_coord_x[ex-1][slot_coord_element];
    475:           const PetscScalar dy = arr_coord_y[ey][slot_coord_next]    - arr_coord_y[ey  ][slot_coord_prev];
    476:           const PetscScalar dz = arr_coord_z[ez][slot_coord_next]    - arr_coord_z[ez  ][slot_coord_prev];
    477:           const PetscScalar B = arr_buoyancy[ez][ey][ex][slot_buoyancy_left];

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

    485:         /* Update z-velocity */
    486:         if (ez > 0) {
    487:           const PetscScalar dx = arr_coord_x[ex][slot_coord_next]    - arr_coord_x[ex  ][slot_coord_prev];
    488:           const PetscScalar dy = arr_coord_y[ey][slot_coord_next]    - arr_coord_y[ey  ][slot_coord_prev];
    489:           const PetscScalar dz = arr_coord_z[ez][slot_coord_element] - arr_coord_z[ez-1][slot_coord_element];
    490:           const PetscScalar B = arr_buoyancy[ez][ey][ex][slot_buoyancy_back];

    492:           arr_velocity[ez][ey][ex][slot_vz_back] += B * ctx->dt * (
    493:               (arr_stress[ez][ey][ex][slot_txz_backright] - arr_stress[ez  ][ey][ex][slot_txz_backleft]) / dx
    494:             + (arr_stress[ez][ey][ex][slot_tyz_backup]    - arr_stress[ez  ][ey][ex][slot_tyz_backdown]) / dy
    495:             + (arr_stress[ez][ey][ex][slot_tzz]           - arr_stress[ez-1][ey][ex][slot_tzz]) / dz);
    496:         }
    497:       }
    498:     }
    499:   }

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

    513: static PetscErrorCode UpdateVelocity(const Ctx *ctx,Vec velocity,Vec stress, Vec buoyancy)
    514: {
    516:   if (ctx->dim == 2) {
    517:     UpdateVelocity_2d(ctx,velocity,stress,buoyancy);
    518:   } else if (ctx->dim == 3) {
    519:     UpdateVelocity_3d(ctx,velocity,stress,buoyancy);
    520:   } else SETERRQ(PetscObjectComm((PetscObject)ctx->dm_velocity),PETSC_ERR_SUP,"Unsupported dim %d",ctx->dim);
    521:   return 0;
    522: }

    524: static PetscErrorCode UpdateStress_2d(const Ctx *ctx,Vec velocity,Vec stress, Vec lame)
    525: {
    526:   Vec               velocity_local,stress_local,lame_local;
    527:   PetscInt          ex,ey,startx,starty,nx,ny;
    528:   PetscInt          slot_coord_next,slot_coord_element,slot_coord_prev;
    529:   PetscInt          slot_vx_left,slot_vy_down,slot_vx_right,slot_vy_up;
    530:   PetscInt          slot_mu_element,slot_lambda_element,slot_mu_downleft;
    531:   PetscInt          slot_txx,slot_tyy,slot_txy_downleft;
    532:   const PetscScalar **arr_coord_x,**arr_coord_y;
    533:   const PetscScalar ***arr_velocity,***arr_lame;
    534:   PetscScalar       ***arr_stress;

    538:   /* Prepare read-write access to stress data */
    539:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_ELEMENT,0,&slot_txx);
    540:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_ELEMENT,1,&slot_tyy);
    541:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_DOWN_LEFT,0,&slot_txy_downleft);
    542:   DMGetLocalVector(ctx->dm_stress,&stress_local);
    543:   DMGlobalToLocal(ctx->dm_stress,stress,INSERT_VALUES,stress_local);
    544:   DMStagVecGetArray(ctx->dm_stress,stress_local,&arr_stress);

    546:   /* Prepare read-only access to velocity data */
    547:   DMStagGetLocationSlot(ctx->dm_velocity,DMSTAG_LEFT,0,&slot_vx_left);
    548:   DMStagGetLocationSlot(ctx->dm_velocity,DMSTAG_RIGHT,0,&slot_vx_right);
    549:   DMStagGetLocationSlot(ctx->dm_velocity,DMSTAG_DOWN,0,&slot_vy_down);
    550:   DMStagGetLocationSlot(ctx->dm_velocity,DMSTAG_UP,0,&slot_vy_up);
    551:   DMGetLocalVector(ctx->dm_velocity,&velocity_local);
    552:   DMGlobalToLocal(ctx->dm_velocity,velocity,INSERT_VALUES,velocity_local);
    553:   DMStagVecGetArrayRead(ctx->dm_velocity,velocity_local,&arr_velocity);

    555:   /* Prepare read-only access to Lame' data */
    556:   DMStagGetLocationSlot(ctx->dm_lame,DMSTAG_ELEMENT,0,&slot_lambda_element);
    557:   DMStagGetLocationSlot(ctx->dm_lame,DMSTAG_ELEMENT,1,&slot_mu_element);
    558:   DMStagGetLocationSlot(ctx->dm_lame,DMSTAG_DOWN_LEFT,0,&slot_mu_downleft);
    559:   DMGetLocalVector(ctx->dm_lame,&lame_local);
    560:   DMGlobalToLocal(ctx->dm_lame,lame,INSERT_VALUES,lame_local);
    561:   DMStagVecGetArrayRead(ctx->dm_lame,lame_local,&arr_lame);

    563:   /* Prepare read-only access to coordinate data */
    564:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity,DMSTAG_LEFT,&slot_coord_prev);
    565:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity,DMSTAG_RIGHT,&slot_coord_next);
    566:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity,DMSTAG_ELEMENT,&slot_coord_element);
    567:   DMStagGetProductCoordinateArrays(ctx->dm_velocity,&arr_coord_x,&arr_coord_y,NULL);

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

    574:       /* Update tau_xx and tau_yy*/
    575:       {
    576:         const PetscScalar dx   = arr_coord_x[ex][slot_coord_next] - arr_coord_x[ex][slot_coord_prev];
    577:         const PetscScalar dy   = arr_coord_y[ey][slot_coord_next] - arr_coord_y[ey][slot_coord_prev];
    578:         const PetscScalar L    = arr_lame[ey][ex][slot_lambda_element];
    579:         const PetscScalar Lp2M = arr_lame[ey][ex][slot_lambda_element] + 2 * arr_lame[ey][ex][slot_mu_element];

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

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

    590:       /* Update tau_xy */
    591:       if (ex > 0 && ey > 0) {
    592:         const PetscScalar dx = arr_coord_x[ex][slot_coord_element] - arr_coord_x[ex-1][slot_coord_element];
    593:         const PetscScalar dy = arr_coord_y[ey][slot_coord_element] - arr_coord_y[ey-1][slot_coord_element];
    594:         const PetscScalar M  = arr_lame[ey][ex][slot_mu_downleft];

    596:         arr_stress[ey][ex][slot_txy_downleft] += M * ctx->dt * (
    597:             (arr_velocity[ey][ex][slot_vx_left] - arr_velocity[ey-1][ex][slot_vx_left]) / dy
    598:           + (arr_velocity[ey][ex][slot_vy_down] - arr_velocity[ey][ex-1][slot_vy_down]) / dx);
    599:       }
    600:     }
    601:   }

    603:   /* Restore all access */
    604:   DMStagVecRestoreArray(ctx->dm_stress,stress_local,&arr_stress);
    605:   DMLocalToGlobal(ctx->dm_stress,stress_local,INSERT_VALUES,stress);
    606:   DMRestoreLocalVector(ctx->dm_stress,&stress_local);
    607:   DMStagVecRestoreArrayRead(ctx->dm_velocity,velocity_local,&arr_velocity);
    608:   DMRestoreLocalVector(ctx->dm_velocity,&velocity_local);
    609:   DMStagVecRestoreArrayRead(ctx->dm_lame,lame_local,&arr_lame);
    610:   DMRestoreLocalVector(ctx->dm_lame,&lame_local);
    611:   DMStagRestoreProductCoordinateArrays(ctx->dm_velocity,&arr_coord_x,&arr_coord_y,NULL);
    612:   return 0;
    613: }

    615: static PetscErrorCode UpdateStress_3d(const Ctx *ctx,Vec velocity,Vec stress, Vec lame)
    616: {
    617:   Vec               velocity_local,stress_local,lame_local;
    618:   PetscInt          ex,ey,ez,startx,starty,startz,nx,ny,nz;
    619:   PetscInt          slot_coord_next,slot_coord_element,slot_coord_prev;
    620:   PetscInt          slot_vx_left,slot_vy_down,slot_vx_right,slot_vy_up,slot_vz_back,slot_vz_front;
    621:   PetscInt          slot_mu_element,slot_lambda_element,slot_mu_downleft,slot_mu_backdown,slot_mu_backleft;
    622:   PetscInt          slot_txx,slot_tyy,slot_tzz,slot_txy_downleft,slot_txz_backleft,slot_tyz_backdown;
    623:   const PetscScalar **arr_coord_x,**arr_coord_y,**arr_coord_z;
    624:   const PetscScalar ****arr_velocity,****arr_lame;
    625:   PetscScalar       ****arr_stress;

    629:   /* Prepare read-write access to stress data */
    630:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_ELEMENT,0,&slot_txx);
    631:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_ELEMENT,1,&slot_tyy);
    632:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_ELEMENT,2,&slot_tzz);
    633:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_DOWN_LEFT,0,&slot_txy_downleft);
    634:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_BACK_LEFT,0,&slot_txz_backleft);
    635:   DMStagGetLocationSlot(ctx->dm_stress,DMSTAG_BACK_DOWN,0,&slot_tyz_backdown);
    636:   DMGetLocalVector(ctx->dm_stress,&stress_local);
    637:   DMGlobalToLocal(ctx->dm_stress,stress,INSERT_VALUES,stress_local);
    638:   DMStagVecGetArray(ctx->dm_stress,stress_local,&arr_stress);

    640:   /* Prepare read-only access to velocity data */
    641:   DMStagGetLocationSlot(ctx->dm_velocity,DMSTAG_LEFT, 0,&slot_vx_left);
    642:   DMStagGetLocationSlot(ctx->dm_velocity,DMSTAG_RIGHT,0,&slot_vx_right);
    643:   DMStagGetLocationSlot(ctx->dm_velocity,DMSTAG_DOWN, 0,&slot_vy_down);
    644:   DMStagGetLocationSlot(ctx->dm_velocity,DMSTAG_UP,   0,&slot_vy_up);
    645:   DMStagGetLocationSlot(ctx->dm_velocity,DMSTAG_BACK, 0,&slot_vz_back);
    646:   DMStagGetLocationSlot(ctx->dm_velocity,DMSTAG_FRONT,0,&slot_vz_front);
    647:   DMGetLocalVector(ctx->dm_velocity,&velocity_local);
    648:   DMGlobalToLocal(ctx->dm_velocity,velocity,INSERT_VALUES,velocity_local);
    649:   DMStagVecGetArrayRead(ctx->dm_velocity,velocity_local,&arr_velocity);

    651:   /* Prepare read-only access to Lame' data */
    652:   DMStagGetLocationSlot(ctx->dm_lame,DMSTAG_ELEMENT,  0,&slot_lambda_element);
    653:   DMStagGetLocationSlot(ctx->dm_lame,DMSTAG_ELEMENT,  1,&slot_mu_element);
    654:   DMStagGetLocationSlot(ctx->dm_lame,DMSTAG_DOWN_LEFT,0,&slot_mu_downleft);
    655:   DMStagGetLocationSlot(ctx->dm_lame,DMSTAG_BACK_LEFT,0,&slot_mu_backleft);
    656:   DMStagGetLocationSlot(ctx->dm_lame,DMSTAG_BACK_DOWN,0,&slot_mu_backdown);
    657:   DMGetLocalVector(ctx->dm_lame,&lame_local);
    658:   DMGlobalToLocal(ctx->dm_lame,lame,INSERT_VALUES,lame_local);
    659:   DMStagVecGetArrayRead(ctx->dm_lame,lame_local,&arr_lame);

    661:   /* Prepare read-only access to coordinate data */
    662:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity,DMSTAG_LEFT,&slot_coord_prev);
    663:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity,DMSTAG_RIGHT,&slot_coord_next);
    664:   DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity,DMSTAG_ELEMENT,&slot_coord_element);
    665:   DMStagGetProductCoordinateArrays(ctx->dm_velocity,&arr_coord_x,&arr_coord_y,&arr_coord_z);

    667:   /* Iterate over the interior of the domain, updating the stresses */
    668:   DMStagGetCorners(ctx->dm_velocity,&startx,&starty,&startz,&nx,&ny,&nz,NULL,NULL,NULL);
    669:   for (ez = startz; ez < startz + nz; ++ez) {
    670:     for (ey = starty; ey < starty + ny; ++ey) {
    671:       for (ex = startx; ex < startx + nx; ++ex) {

    673:         /* Update tau_xx, tau_yy, and tau_zz*/
    674:         {
    675:           const PetscScalar dx   = arr_coord_x[ex][slot_coord_next] - arr_coord_x[ex][slot_coord_prev];
    676:           const PetscScalar dy   = arr_coord_y[ey][slot_coord_next] - arr_coord_y[ey][slot_coord_prev];
    677:           const PetscScalar dz   = arr_coord_z[ez][slot_coord_next] - arr_coord_z[ez][slot_coord_prev];
    678:           const PetscScalar L    = arr_lame[ez][ey][ex][slot_lambda_element];
    679:           const PetscScalar Lp2M = arr_lame[ez][ey][ex][slot_lambda_element] + 2 * arr_lame[ez][ey][ex][slot_mu_element];

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

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

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

    697:         /* Update tau_xy, tau_xz, tau_yz */
    698:         {
    699:           PetscScalar dx,dy,dz;

    701:           if (ex > 0)  dx = arr_coord_x[ex][slot_coord_element] - arr_coord_x[ex-1][slot_coord_element];
    702:           if (ey > 0)  dy = arr_coord_y[ey][slot_coord_element] - arr_coord_y[ey-1][slot_coord_element];
    703:           if (ez > 0)  dz = arr_coord_z[ez][slot_coord_element] - arr_coord_z[ez-1][slot_coord_element];

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

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

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

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

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

    726:             arr_stress[ez][ey][ex][slot_tyz_backdown] += M * ctx->dt * (
    727:                 (arr_velocity[ez][ey][ex][slot_vy_down] - arr_velocity[ez-1][ey][ex][slot_vy_down]) / dz
    728:               + (arr_velocity[ez][ey][ex][slot_vz_back] - arr_velocity[ez][ey-1][ex][slot_vz_back]) / dy);
    729:           }
    730:         }
    731:       }
    732:     }
    733:   }

    735:   /* Restore all access */
    736:   DMStagVecRestoreArray(ctx->dm_stress,stress_local,&arr_stress);
    737:   DMLocalToGlobal(ctx->dm_stress,stress_local,INSERT_VALUES,stress);
    738:   DMRestoreLocalVector(ctx->dm_stress,&stress_local);
    739:   DMStagVecRestoreArrayRead(ctx->dm_velocity,velocity_local,&arr_velocity);
    740:   DMRestoreLocalVector(ctx->dm_velocity,&velocity_local);
    741:   DMStagVecRestoreArrayRead(ctx->dm_lame,lame_local,&arr_lame);
    742:   DMRestoreLocalVector(ctx->dm_lame,&lame_local);
    743:   DMStagRestoreProductCoordinateArrays(ctx->dm_velocity,&arr_coord_x,&arr_coord_y,&arr_coord_z);
    744:   return 0;
    745: }

    747: static PetscErrorCode UpdateStress(const Ctx *ctx,Vec velocity,Vec stress, Vec lame)
    748: {
    750:   if (ctx->dim == 2) {
    751:     UpdateStress_2d(ctx,velocity,stress,lame);
    752:   } else if (ctx->dim ==3) {
    753:     UpdateStress_3d(ctx,velocity,stress,lame);
    754:   }
    755:   return 0;
    756: }

    758: static PetscErrorCode DumpStress(const Ctx *ctx,Vec stress,PetscInt timestep)
    759: {
    760:   DM             da_normal,da_shear = NULL;
    761:   Vec            vec_normal,vec_shear = NULL;

    765:     DMStagVecSplitToDMDA(ctx->dm_stress,stress,DMSTAG_ELEMENT,-ctx->dim,&da_normal,&vec_normal);
    766:     PetscObjectSetName((PetscObject)vec_normal,"normal stresses");

    768:     /* Dump element-based fields to a .vtr file */
    769:     {
    770:       PetscViewer viewer;
    771:       char        filename[PETSC_MAX_PATH_LEN];

    773:       PetscSNPrintf(filename,sizeof(filename),"ex6_stress_normal_%.4D.vtr",timestep);
    774:       PetscViewerVTKOpen(PetscObjectComm((PetscObject)da_normal),filename,FILE_MODE_WRITE,&viewer);
    775:       VecView(vec_normal,viewer);
    776:       PetscViewerDestroy(&viewer);
    777:       PetscPrintf(PETSC_COMM_WORLD,"Created %s\n",filename);
    778:     }

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

    785:     {
    786:       PetscViewer viewer;
    787:       char        filename[PETSC_MAX_PATH_LEN];

    789:       PetscSNPrintf(filename,sizeof(filename),"ex6_stress_shear_%.4D.vtr",timestep);
    790:       PetscViewerVTKOpen(PetscObjectComm((PetscObject)da_normal),filename,FILE_MODE_WRITE,&viewer);
    791:       VecView(vec_shear,viewer);
    792:       PetscViewerDestroy(&viewer);
    793:       PetscPrintf(PETSC_COMM_WORLD,"Created %s\n",filename);
    794:     }
    795:   }

    797:   /* Destroy DMDAs and Vecs */
    798:   DMDestroy(&da_normal);
    799:   DMDestroy(&da_shear);
    800:   VecDestroy(&vec_normal);
    801:   VecDestroy(&vec_shear);
    802:   return 0;
    803: }

    805: static PetscErrorCode DumpVelocity(const Ctx *ctx,Vec velocity,PetscInt timestep)
    806: {
    807:   DM             dmVelAvg;
    808:   Vec            velAvg;
    809:   DM             daVelAvg;
    810:   Vec            vecVelAvg;
    811:   Vec            velocity_local;
    812:   PetscInt       ex,ey,ez,startx,starty,startz,nx,ny,nz;

    816:   if (ctx->dim == 2) {
    817:     DMStagCreateCompatibleDMStag(ctx->dm_velocity,0,0,2,0,&dmVelAvg); /* 2 dof per element */
    818:   } else if (ctx->dim == 3) {
    819:     DMStagCreateCompatibleDMStag(ctx->dm_velocity,0,0,0,3,&dmVelAvg); /* 3 dof per element */
    820:   } else SETERRQ(PetscObjectComm((PetscObject)ctx->dm_velocity),PETSC_ERR_SUP,"Unsupported dim %d",ctx->dim);
    821:   DMSetUp(dmVelAvg);
    822:   DMStagSetUniformCoordinatesProduct(dmVelAvg,ctx->xmin,ctx->xmax,ctx->ymin,ctx->ymax,ctx->zmin,ctx->zmax);
    823:   DMCreateGlobalVector(dmVelAvg,&velAvg);
    824:   DMGetLocalVector(ctx->dm_velocity,&velocity_local);
    825:   DMGlobalToLocal(ctx->dm_velocity,velocity,INSERT_VALUES,velocity_local);
    826:   DMStagGetCorners(dmVelAvg,&startx,&starty,&startz,&nx,&ny,&nz,NULL,NULL,NULL);
    827:   if (ctx->dim == 2) {
    828:     for (ey = starty; ey<starty+ny; ++ey) {
    829:       for (ex = startx; ex<startx+nx; ++ex) {
    830:         DMStagStencil from[4],to[2];
    831:         PetscScalar   valFrom[4],valTo[2];

    833:         from[0].i = ex; from[0].j = ey; from[0].loc = DMSTAG_UP;    from[0].c = 0;
    834:         from[1].i = ex; from[1].j = ey; from[1].loc = DMSTAG_DOWN;  from[1].c = 0;
    835:         from[2].i = ex; from[2].j = ey; from[2].loc = DMSTAG_LEFT;  from[2].c = 0;
    836:         from[3].i = ex; from[3].j = ey; from[3].loc = DMSTAG_RIGHT; from[3].c = 0;
    837:         DMStagVecGetValuesStencil(ctx->dm_velocity,velocity_local,4,from,valFrom);
    838:         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]);
    839:         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]);
    840:         DMStagVecSetValuesStencil(dmVelAvg,velAvg,2,to,valTo,INSERT_VALUES);
    841:       }
    842:     }
    843:   } else if (ctx->dim == 3) {
    844:     for (ez = startz; ez<startz+nz; ++ez) {
    845:       for (ey = starty; ey<starty+ny; ++ey) {
    846:         for (ex = startx; ex<startx+nx; ++ex) {
    847:           DMStagStencil from[6],to[3];
    848:           PetscScalar   valFrom[6],valTo[3];

    850:           from[0].i = ex; from[0].j = ey; from[0].k = ez; from[0].loc = DMSTAG_UP;    from[0].c = 0;
    851:           from[1].i = ex; from[1].j = ey; from[1].k = ez; from[1].loc = DMSTAG_DOWN;  from[1].c = 0;
    852:           from[2].i = ex; from[2].j = ey; from[2].k = ez; from[2].loc = DMSTAG_LEFT;  from[2].c = 0;
    853:           from[3].i = ex; from[3].j = ey; from[3].k = ez; from[3].loc = DMSTAG_RIGHT; from[3].c = 0;
    854:           from[4].i = ex; from[4].j = ey; from[4].k = ez; from[4].loc = DMSTAG_FRONT; from[4].c = 0;
    855:           from[5].i = ex; from[5].j = ey; from[5].k = ez; from[5].loc = DMSTAG_BACK;  from[5].c = 0;
    856:           DMStagVecGetValuesStencil(ctx->dm_velocity,velocity_local,6,from,valFrom);
    857:           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]);
    858:           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]);
    859:           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]);
    860:           DMStagVecSetValuesStencil(dmVelAvg,velAvg,3,to,valTo,INSERT_VALUES);
    861:         }
    862:       }
    863:     }
    864:   } else SETERRQ(PetscObjectComm((PetscObject)ctx->dm_velocity),PETSC_ERR_SUP,"Unsupported dim %d",ctx->dim);
    865:   VecAssemblyBegin(velAvg);
    866:   VecAssemblyEnd(velAvg);
    867:   DMRestoreLocalVector(ctx->dm_velocity,&velocity_local);

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

    872:   /* Dump element-based fields to a .vtr file */
    873:   {
    874:     PetscViewer viewer;
    875:     char        filename[PETSC_MAX_PATH_LEN];

    877:     PetscSNPrintf(filename,sizeof(filename),"ex6_velavg_%.4D.vtr",timestep);
    878:     PetscViewerVTKOpen(PetscObjectComm((PetscObject)daVelAvg),filename,FILE_MODE_WRITE,&viewer);
    879:     VecView(vecVelAvg,viewer);
    880:     PetscViewerDestroy(&viewer);
    881:     PetscPrintf(PETSC_COMM_WORLD,"Created %s\n",filename);
    882:   }

    884:   /* Destroy DMDAs and Vecs */
    885:   VecDestroy(&vecVelAvg);
    886:   DMDestroy(&daVelAvg);
    887:   VecDestroy(&velAvg);
    888:   DMDestroy(&dmVelAvg);
    889:   return 0;
    890: }

    892: /*TEST

    894:    test:
    895:       suffix: 1
    896:       requires: !complex
    897:       nsize: 1
    898:       args: -stag_grid_x 12 -stag_grid_y 7 -nsteps 4 -dump_output 0

    900:    test:
    901:       suffix: 2
    902:       requires: !complex
    903:       nsize: 9
    904:       args: -stag_grid_x 12 -stag_grid_y 15 -nsteps 3 -dump_output 0

    906:    test:
    907:       suffix: 3
    908:       requires: !complex
    909:       nsize: 1
    910:       args: -stag_grid_x 12 -stag_grid_y 7 -stag_grid_z 11 -nsteps 2 -dim 3 -dump_output 0

    912:    test:
    913:       suffix: 4
    914:       requires: !complex
    915:       nsize: 12
    916:       args: -stag_grid_x 12 -stag_grid_y 15 -stag_grid_z 8 -nsteps 3 -dim 3 -dump_output 0

    918: TEST*/