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";
7
8 /*
9 Reference:
10
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 }
21
22 This example uses y in 2 dimensions, where the paper uses z.
23
24 This example uses the dual grid of the one pictured in Fig. 1. of the paper, so
25 that velocities are on face boundaries, shear stresses are defined on
26 vertices(2D) or edges(3D), and normal stresses are defined on elements.
27
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
32 This example treats the boundaries naively (by leaving ~zero velocity and
33 stress there).
34 */
35
36 #include <petscdmstag.h>
37 #include <petscts.h>
38
39 /* A struct defining the parameters of the problem */
40 typedef struct {
41 DM dm_velocity, dm_stress;
42 DM dm_buoyancy, dm_lame;
43 Vec buoyancy, lame; /* Global, but for efficiency one could store local vectors */
44 PetscInt dim; /* 2 or 3 */
45 PetscScalar rho, lambda, mu; /* constant material properties */
46 PetscReal xmin, xmax, ymin, ymax, zmin, zmax;
47 PetscReal dt;
48 PetscInt timesteps;
49 PetscBool dump_output;
50 } Ctx;
51
52 static PetscErrorCode CreateLame(Ctx *);
53 static PetscErrorCode ForceStress(const Ctx *, Vec, PetscReal);
54 static PetscErrorCode DumpVelocity(const Ctx *, Vec, PetscInt);
55 static PetscErrorCode DumpStress(const Ctx *, Vec, PetscInt);
56 static PetscErrorCode UpdateVelocity(const Ctx *, Vec, Vec, Vec);
57 static PetscErrorCode UpdateStress(const Ctx *, Vec, Vec, Vec);
58
main(int argc,char * argv[])59 int main(int argc, char *argv[])
60 {
61 Ctx ctx;
62 Vec velocity, stress;
63 PetscInt timestep;
64
65 /* Initialize PETSc */
66 PetscFunctionBeginUser;
67 PetscCall(PetscInitialize(&argc, &argv, 0, help));
68
69 /* Populate application context */
70 ctx.dim = 2;
71 ctx.rho = 1.0;
72 ctx.lambda = 1.0;
73 ctx.mu = 1.0;
74 ctx.xmin = 0.0;
75 ctx.xmax = 1.0;
76 ctx.ymin = 0.0;
77 ctx.ymax = 1.0;
78 ctx.zmin = 0.0;
79 ctx.zmax = 1.0;
80 ctx.dt = 0.001;
81 ctx.timesteps = 100;
82 ctx.dump_output = PETSC_TRUE;
83
84 /* Update context from command line options */
85 PetscCall(PetscOptionsGetInt(NULL, NULL, "-dim", &ctx.dim, NULL));
86 PetscCall(PetscOptionsGetReal(NULL, NULL, "-dt", &ctx.dt, NULL));
87 PetscCall(PetscOptionsGetInt(NULL, NULL, "-nsteps", &ctx.timesteps, NULL));
88 PetscCall(PetscOptionsGetBool(NULL, NULL, "-dump_output", &ctx.dump_output, NULL));
89
90 /* Create a DMStag, with uniform coordinates, for the velocities */
91 {
92 PetscInt dof0, dof1, dof2, dof3;
93 const PetscInt stencilWidth = 1;
94
95 switch (ctx.dim) {
96 case 2:
97 dof0 = 0;
98 dof1 = 1;
99 dof2 = 0; /* 1 dof per cell boundary */
100 PetscCall(DMStagCreate2d(PETSC_COMM_WORLD, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, 100, 100, PETSC_DECIDE, PETSC_DECIDE, dof0, dof1, dof2, DMSTAG_STENCIL_BOX, stencilWidth, NULL, NULL, &ctx.dm_velocity));
101 break;
102 case 3:
103 dof0 = 0;
104 dof1 = 0;
105 dof2 = 1;
106 dof3 = 0; /* 1 dof per cell boundary */
107 PetscCall(DMStagCreate3d(PETSC_COMM_WORLD, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, 30, 30, 30, PETSC_DECIDE, PETSC_DECIDE, PETSC_DECIDE, dof0, dof1, dof2, dof3, DMSTAG_STENCIL_BOX, stencilWidth, NULL, NULL, NULL, &ctx.dm_velocity));
108 break;
109 default:
110 SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_SUP, "Not Implemented for dimension %" PetscInt_FMT, ctx.dim);
111 }
112 }
113 PetscCall(DMSetFromOptions(ctx.dm_velocity)); /* Options control velocity DM */
114 PetscCall(DMSetUp(ctx.dm_velocity));
115 PetscCall(DMStagSetUniformCoordinatesProduct(ctx.dm_velocity, ctx.xmin, ctx.xmax, ctx.ymin, ctx.ymax, ctx.zmin, ctx.zmax));
116
117 /* Create a second, compatible DMStag for the stresses */
118 switch (ctx.dim) {
119 case 2:
120 /* One shear stress component on element corners, two shear stress components on elements */
121 PetscCall(DMStagCreateCompatibleDMStag(ctx.dm_velocity, 1, 0, 2, 0, &ctx.dm_stress));
122 break;
123 case 3:
124 /* One shear stress component on element edges, three shear stress components on elements */
125 PetscCall(DMStagCreateCompatibleDMStag(ctx.dm_velocity, 0, 1, 0, 3, &ctx.dm_stress));
126 break;
127 default:
128 SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_SUP, "Not Implemented for dimension %" PetscInt_FMT, ctx.dim);
129 }
130 PetscCall(DMSetUp(ctx.dm_stress));
131 PetscCall(DMStagSetUniformCoordinatesProduct(ctx.dm_stress, ctx.xmin, ctx.xmax, ctx.ymin, ctx.ymax, ctx.zmin, ctx.zmax));
132
133 /* Create two additional DMStag objects for the buoyancy and Lame parameters */
134 switch (ctx.dim) {
135 case 2:
136 /* buoyancy on element boundaries (edges) */
137 PetscCall(DMStagCreateCompatibleDMStag(ctx.dm_velocity, 0, 1, 0, 0, &ctx.dm_buoyancy));
138 break;
139 case 3:
140 /* buoyancy on element boundaries (faces) */
141 PetscCall(DMStagCreateCompatibleDMStag(ctx.dm_velocity, 0, 0, 1, 0, &ctx.dm_buoyancy));
142 break;
143 default:
144 SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_SUP, "Not Implemented for dimension %" PetscInt_FMT, ctx.dim);
145 }
146 PetscCall(DMSetUp(ctx.dm_buoyancy));
147
148 switch (ctx.dim) {
149 case 2:
150 /* mu and lambda + 2*mu on element centers, mu on corners */
151 PetscCall(DMStagCreateCompatibleDMStag(ctx.dm_velocity, 1, 0, 2, 0, &ctx.dm_lame));
152 break;
153 case 3:
154 /* mu and lambda + 2*mu on element centers, mu on edges */
155 PetscCall(DMStagCreateCompatibleDMStag(ctx.dm_velocity, 0, 1, 0, 2, &ctx.dm_lame));
156 break;
157 default:
158 SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_SUP, "Not Implemented for dimension %" PetscInt_FMT, ctx.dim);
159 }
160 PetscCall(DMSetUp(ctx.dm_lame));
161
162 /* Print out some info */
163 {
164 PetscInt N[3];
165 PetscScalar dx, Vp;
166
167 PetscCall(DMStagGetGlobalSizes(ctx.dm_velocity, &N[0], &N[1], &N[2]));
168 dx = (ctx.xmax - ctx.xmin) / N[0];
169 Vp = PetscSqrtScalar((ctx.lambda + 2 * ctx.mu) / ctx.rho);
170 if (ctx.dim == 2) {
171 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Using a %" PetscInt_FMT " x %" PetscInt_FMT " mesh\n", N[0], N[1]));
172 } else if (ctx.dim == 3) {
173 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Using a %" PetscInt_FMT " x %" PetscInt_FMT " x %" PetscInt_FMT " mesh\n", N[0], N[1], N[2]));
174 }
175 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "dx: %g\n", (double)PetscRealPart(dx)));
176 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "dt: %g\n", (double)ctx.dt));
177 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "P-wave velocity: %g\n", (double)PetscRealPart(PetscSqrtScalar((ctx.lambda + 2 * ctx.mu) / ctx.rho))));
178 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "V_p dt / dx: %g\n", (double)PetscRealPart(Vp * ctx.dt / dx)));
179 }
180
181 /* Populate the coefficient arrays */
182 PetscCall(CreateLame(&ctx));
183 PetscCall(DMCreateGlobalVector(ctx.dm_buoyancy, &ctx.buoyancy));
184 PetscCall(VecSet(ctx.buoyancy, 1.0 / ctx.rho));
185
186 /* Create vectors to store the system state */
187 PetscCall(DMCreateGlobalVector(ctx.dm_velocity, &velocity));
188 PetscCall(DMCreateGlobalVector(ctx.dm_stress, &stress));
189
190 /* Initial State */
191 PetscCall(VecSet(velocity, 0.0));
192 PetscCall(VecSet(stress, 0.0));
193 PetscCall(ForceStress(&ctx, stress, 0.0));
194 if (ctx.dump_output) {
195 PetscCall(DumpVelocity(&ctx, velocity, 0));
196 PetscCall(DumpStress(&ctx, stress, 0));
197 }
198
199 /* Time Loop */
200 for (timestep = 1; timestep <= ctx.timesteps; ++timestep) {
201 const PetscReal t = timestep * ctx.dt;
202
203 PetscCall(UpdateVelocity(&ctx, velocity, stress, ctx.buoyancy));
204 PetscCall(UpdateStress(&ctx, velocity, stress, ctx.lame));
205 PetscCall(ForceStress(&ctx, stress, t));
206 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Timestep %" PetscInt_FMT ", t = %g\n", timestep, (double)t));
207 if (ctx.dump_output) {
208 PetscCall(DumpVelocity(&ctx, velocity, timestep));
209 PetscCall(DumpStress(&ctx, stress, timestep));
210 }
211 }
212
213 /* Clean up and finalize PETSc */
214 PetscCall(VecDestroy(&velocity));
215 PetscCall(VecDestroy(&stress));
216 PetscCall(VecDestroy(&ctx.lame));
217 PetscCall(VecDestroy(&ctx.buoyancy));
218 PetscCall(DMDestroy(&ctx.dm_velocity));
219 PetscCall(DMDestroy(&ctx.dm_stress));
220 PetscCall(DMDestroy(&ctx.dm_buoyancy));
221 PetscCall(DMDestroy(&ctx.dm_lame));
222 PetscCall(PetscFinalize());
223 return 0;
224 }
225
CreateLame(Ctx * ctx)226 static PetscErrorCode CreateLame(Ctx *ctx)
227 {
228 PetscInt N[3], ex, ey, ez, startx, starty, startz, nx, ny, nz, extrax, extray, extraz;
229
230 PetscFunctionBeginUser;
231 PetscCall(DMCreateGlobalVector(ctx->dm_lame, &ctx->lame));
232 PetscCall(DMStagGetGlobalSizes(ctx->dm_lame, &N[0], &N[1], &N[2]));
233 PetscCall(DMStagGetCorners(ctx->dm_buoyancy, &startx, &starty, &startz, &nx, &ny, &nz, &extrax, &extray, &extraz));
234 if (ctx->dim == 2) {
235 /* Element values */
236 for (ey = starty; ey < starty + ny; ++ey) {
237 for (ex = startx; ex < startx + nx; ++ex) {
238 DMStagStencil pos;
239
240 pos.i = ex;
241 pos.j = ey;
242 pos.c = 0;
243 pos.loc = DMSTAG_ELEMENT;
244 PetscCall(DMStagVecSetValuesStencil(ctx->dm_lame, ctx->lame, 1, &pos, &ctx->lambda, INSERT_VALUES));
245 pos.i = ex;
246 pos.j = ey;
247 pos.c = 1;
248 pos.loc = DMSTAG_ELEMENT;
249 PetscCall(DMStagVecSetValuesStencil(ctx->dm_lame, ctx->lame, 1, &pos, &ctx->mu, INSERT_VALUES));
250 }
251 }
252 /* Vertex Values */
253 for (ey = starty; ey < starty + ny + extray; ++ey) {
254 for (ex = startx; ex < startx + nx + extrax; ++ex) {
255 DMStagStencil pos;
256
257 pos.i = ex;
258 pos.j = ey;
259 pos.c = 0;
260 pos.loc = DMSTAG_DOWN_LEFT;
261 PetscCall(DMStagVecSetValuesStencil(ctx->dm_lame, ctx->lame, 1, &pos, &ctx->mu, INSERT_VALUES));
262 }
263 }
264 } else if (ctx->dim == 3) {
265 /* Element values */
266 for (ez = startz; ez < startz + nz; ++ez) {
267 for (ey = starty; ey < starty + ny; ++ey) {
268 for (ex = startx; ex < startx + nx; ++ex) {
269 DMStagStencil pos;
270
271 pos.i = ex;
272 pos.j = ey;
273 pos.k = ez;
274 pos.c = 0;
275 pos.loc = DMSTAG_ELEMENT;
276 PetscCall(DMStagVecSetValuesStencil(ctx->dm_lame, ctx->lame, 1, &pos, &ctx->lambda, INSERT_VALUES));
277 pos.i = ex;
278 pos.j = ey;
279 pos.k = ez;
280 pos.c = 1;
281 pos.loc = DMSTAG_ELEMENT;
282 PetscCall(DMStagVecSetValuesStencil(ctx->dm_lame, ctx->lame, 1, &pos, &ctx->mu, INSERT_VALUES));
283 }
284 }
285 }
286 /* Edge Values */
287 for (ez = startz; ez < startz + nz + extraz; ++ez) {
288 for (ey = starty; ey < starty + ny + extray; ++ey) {
289 for (ex = startx; ex < startx + nx + extrax; ++ex) {
290 DMStagStencil pos;
291
292 if (ex < N[0]) {
293 pos.i = ex;
294 pos.j = ey;
295 pos.k = ez;
296 pos.c = 0;
297 pos.loc = DMSTAG_BACK_DOWN;
298 PetscCall(DMStagVecSetValuesStencil(ctx->dm_lame, ctx->lame, 1, &pos, &ctx->mu, INSERT_VALUES));
299 }
300 if (ey < N[1]) {
301 pos.i = ex;
302 pos.j = ey;
303 pos.k = ez;
304 pos.c = 0;
305 pos.loc = DMSTAG_BACK_LEFT;
306 PetscCall(DMStagVecSetValuesStencil(ctx->dm_lame, ctx->lame, 1, &pos, &ctx->mu, INSERT_VALUES));
307 }
308 if (ez < N[2]) {
309 pos.i = ex;
310 pos.j = ey;
311 pos.k = ez;
312 pos.c = 0;
313 pos.loc = DMSTAG_DOWN_LEFT;
314 PetscCall(DMStagVecSetValuesStencil(ctx->dm_lame, ctx->lame, 1, &pos, &ctx->mu, INSERT_VALUES));
315 }
316 }
317 }
318 }
319 } else SETERRQ(PetscObjectComm((PetscObject)ctx->dm_velocity), PETSC_ERR_SUP, "Unsupported dim %" PetscInt_FMT, ctx->dim);
320 PetscCall(VecAssemblyBegin(ctx->lame));
321 PetscCall(VecAssemblyEnd(ctx->lame));
322 PetscFunctionReturn(PETSC_SUCCESS);
323 }
324
ForceStress(const Ctx * ctx,Vec stress,PetscReal t)325 static PetscErrorCode ForceStress(const Ctx *ctx, Vec stress, PetscReal t)
326 {
327 PetscInt start[3], n[3], N[3];
328 DMStagStencil pos;
329 PetscBool this_rank;
330 const PetscScalar val = PetscExpReal(-100.0 * t);
331
332 PetscFunctionBeginUser;
333 PetscCall(DMStagGetGlobalSizes(ctx->dm_stress, &N[0], &N[1], &N[2]));
334 PetscCall(DMStagGetCorners(ctx->dm_stress, &start[0], &start[1], &start[2], &n[0], &n[1], &n[2], NULL, NULL, NULL));
335
336 /* Normal stresses at a single point */
337 this_rank = (PetscBool)(start[0] <= N[0] / 2 && N[0] / 2 <= start[0] + n[0]);
338 this_rank = (PetscBool)(this_rank && start[1] <= N[1] / 2 && N[1] / 2 <= start[1] + n[1]);
339 if (ctx->dim == 3) this_rank = (PetscBool)(this_rank && start[2] <= N[2] / 2 && N[2] / 2 <= start[2] + n[2]);
340 if (this_rank) {
341 /* Note integer division to pick element near the center */
342 pos.i = N[0] / 2;
343 pos.j = N[1] / 2;
344 pos.k = N[2] / 2;
345 pos.c = 0;
346 pos.loc = DMSTAG_ELEMENT;
347 PetscCall(DMStagVecSetValuesStencil(ctx->dm_stress, stress, 1, &pos, &val, INSERT_VALUES));
348 pos.i = N[0] / 2;
349 pos.j = N[1] / 2;
350 pos.k = N[2] / 2;
351 pos.c = 1;
352 pos.loc = DMSTAG_ELEMENT;
353 PetscCall(DMStagVecSetValuesStencil(ctx->dm_stress, stress, 1, &pos, &val, INSERT_VALUES));
354 if (ctx->dim == 3) {
355 pos.i = N[0] / 2;
356 pos.j = N[1] / 2;
357 pos.k = N[2] / 2;
358 pos.c = 2;
359 pos.loc = DMSTAG_ELEMENT;
360 PetscCall(DMStagVecSetValuesStencil(ctx->dm_stress, stress, 1, &pos, &val, INSERT_VALUES));
361 }
362 }
363
364 PetscCall(VecAssemblyBegin(stress));
365 PetscCall(VecAssemblyEnd(stress));
366 PetscFunctionReturn(PETSC_SUCCESS);
367 }
368
UpdateVelocity_2d(const Ctx * ctx,Vec velocity,Vec stress,Vec buoyancy)369 static PetscErrorCode UpdateVelocity_2d(const Ctx *ctx, Vec velocity, Vec stress, Vec buoyancy)
370 {
371 Vec velocity_local, stress_local, buoyancy_local;
372 PetscInt ex, ey, startx, starty, nx, ny;
373 PetscInt slot_coord_next, slot_coord_element, slot_coord_prev;
374 PetscInt slot_vx_left, slot_vy_down, slot_buoyancy_down, slot_buoyancy_left;
375 PetscInt slot_txx, slot_tyy, slot_txy_downleft, slot_txy_downright, slot_txy_upleft;
376 const PetscScalar **arr_coord_x, **arr_coord_y;
377 const PetscScalar ***arr_stress, ***arr_buoyancy;
378 PetscScalar ***arr_velocity;
379
380 PetscFunctionBeginUser;
381 /* Prepare direct access to buoyancy data */
382 PetscCall(DMStagGetLocationSlot(ctx->dm_buoyancy, DMSTAG_LEFT, 0, &slot_buoyancy_left));
383 PetscCall(DMStagGetLocationSlot(ctx->dm_buoyancy, DMSTAG_DOWN, 0, &slot_buoyancy_down));
384 PetscCall(DMGetLocalVector(ctx->dm_buoyancy, &buoyancy_local));
385 PetscCall(DMGlobalToLocal(ctx->dm_buoyancy, buoyancy, INSERT_VALUES, buoyancy_local));
386 PetscCall(DMStagVecGetArrayRead(ctx->dm_buoyancy, buoyancy_local, (void *)&arr_buoyancy));
387
388 /* Prepare read-only access to stress data */
389 PetscCall(DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_ELEMENT, 0, &slot_txx));
390 PetscCall(DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_ELEMENT, 1, &slot_tyy));
391 PetscCall(DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_UP_LEFT, 0, &slot_txy_upleft));
392 PetscCall(DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_DOWN_LEFT, 0, &slot_txy_downleft));
393 PetscCall(DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_DOWN_RIGHT, 0, &slot_txy_downright));
394 PetscCall(DMGetLocalVector(ctx->dm_stress, &stress_local));
395 PetscCall(DMGlobalToLocal(ctx->dm_stress, stress, INSERT_VALUES, stress_local));
396 PetscCall(DMStagVecGetArrayRead(ctx->dm_stress, stress_local, (void *)&arr_stress));
397
398 /* Prepare read-write access to velocity data */
399 PetscCall(DMStagGetLocationSlot(ctx->dm_velocity, DMSTAG_LEFT, 0, &slot_vx_left));
400 PetscCall(DMStagGetLocationSlot(ctx->dm_velocity, DMSTAG_DOWN, 0, &slot_vy_down));
401 PetscCall(DMGetLocalVector(ctx->dm_velocity, &velocity_local));
402 PetscCall(DMGlobalToLocal(ctx->dm_velocity, velocity, INSERT_VALUES, velocity_local));
403 PetscCall(DMStagVecGetArray(ctx->dm_velocity, velocity_local, &arr_velocity));
404
405 /* Prepare read-only access to coordinate data */
406 PetscCall(DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_LEFT, &slot_coord_prev));
407 PetscCall(DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_RIGHT, &slot_coord_next));
408 PetscCall(DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_ELEMENT, &slot_coord_element));
409 PetscCall(DMStagGetProductCoordinateArrays(ctx->dm_velocity, (void *)&arr_coord_x, (void *)&arr_coord_y, NULL));
410
411 /* Iterate over interior of the domain, updating the velocities */
412 PetscCall(DMStagGetCorners(ctx->dm_velocity, &startx, &starty, NULL, &nx, &ny, NULL, NULL, NULL, NULL));
413 for (ey = starty; ey < starty + ny; ++ey) {
414 for (ex = startx; ex < startx + nx; ++ex) {
415 /* Update y-velocity */
416 if (ey > 0) {
417 const PetscScalar dx = arr_coord_x[ex][slot_coord_next] - arr_coord_x[ex][slot_coord_prev];
418 const PetscScalar dy = arr_coord_y[ey][slot_coord_element] - arr_coord_y[ey - 1][slot_coord_element];
419 const PetscScalar B = arr_buoyancy[ey][ex][slot_buoyancy_down];
420
421 arr_velocity[ey][ex][slot_vy_down] += B * ctx->dt * ((arr_stress[ey][ex][slot_txy_downright] - arr_stress[ey][ex][slot_txy_downleft]) / dx + (arr_stress[ey][ex][slot_tyy] - arr_stress[ey - 1][ex][slot_tyy]) / dy);
422 }
423
424 /* Update x-velocity */
425 if (ex > 0) {
426 const PetscScalar dx = arr_coord_x[ex][slot_coord_element] - arr_coord_x[ex - 1][slot_coord_element];
427 const PetscScalar dy = arr_coord_y[ey][slot_coord_next] - arr_coord_y[ey][slot_coord_prev];
428 const PetscScalar B = arr_buoyancy[ey][ex][slot_buoyancy_left];
429
430 arr_velocity[ey][ex][slot_vx_left] += B * ctx->dt * ((arr_stress[ey][ex][slot_txx] - arr_stress[ey][ex - 1][slot_txx]) / dx + (arr_stress[ey][ex][slot_txy_upleft] - arr_stress[ey][ex][slot_txy_downleft]) / dy);
431 }
432 }
433 }
434
435 /* Restore all access */
436 PetscCall(DMStagVecRestoreArrayRead(ctx->dm_buoyancy, buoyancy_local, (void *)&arr_buoyancy));
437 PetscCall(DMRestoreLocalVector(ctx->dm_buoyancy, &buoyancy_local));
438 PetscCall(DMStagVecRestoreArrayRead(ctx->dm_stress, stress_local, (void *)&arr_stress));
439 PetscCall(DMRestoreLocalVector(ctx->dm_stress, &stress_local));
440 PetscCall(DMStagVecRestoreArray(ctx->dm_velocity, velocity_local, &arr_velocity));
441 PetscCall(DMLocalToGlobal(ctx->dm_velocity, velocity_local, INSERT_VALUES, velocity));
442 PetscCall(DMRestoreLocalVector(ctx->dm_velocity, &velocity_local));
443 PetscCall(DMStagRestoreProductCoordinateArrays(ctx->dm_velocity, (void *)&arr_coord_x, (void *)&arr_coord_y, NULL));
444 PetscFunctionReturn(PETSC_SUCCESS);
445 }
446
UpdateVelocity_3d(const Ctx * ctx,Vec velocity,Vec stress,Vec buoyancy)447 static PetscErrorCode UpdateVelocity_3d(const Ctx *ctx, Vec velocity, Vec stress, Vec buoyancy)
448 {
449 Vec velocity_local, stress_local, buoyancy_local;
450 PetscInt ex, ey, ez, startx, starty, startz, nx, ny, nz;
451 PetscInt slot_coord_next, slot_coord_element, slot_coord_prev;
452 PetscInt slot_vx_left, slot_vy_down, slot_vz_back, slot_buoyancy_down, slot_buoyancy_left, slot_buoyancy_back;
453 PetscInt slot_txx, slot_tyy, slot_tzz;
454 PetscInt slot_txy_downleft, slot_txy_downright, slot_txy_upleft;
455 PetscInt slot_txz_backleft, slot_txz_backright, slot_txz_frontleft;
456 PetscInt slot_tyz_backdown, slot_tyz_frontdown, slot_tyz_backup;
457 const PetscScalar **arr_coord_x, **arr_coord_y, **arr_coord_z;
458 const PetscScalar ****arr_stress, ****arr_buoyancy;
459 PetscScalar ****arr_velocity;
460
461 PetscFunctionBeginUser;
462 /* Prepare direct access to buoyancy data */
463 PetscCall(DMStagGetLocationSlot(ctx->dm_buoyancy, DMSTAG_LEFT, 0, &slot_buoyancy_left));
464 PetscCall(DMStagGetLocationSlot(ctx->dm_buoyancy, DMSTAG_DOWN, 0, &slot_buoyancy_down));
465 PetscCall(DMStagGetLocationSlot(ctx->dm_buoyancy, DMSTAG_BACK, 0, &slot_buoyancy_back));
466 PetscCall(DMGetLocalVector(ctx->dm_buoyancy, &buoyancy_local));
467 PetscCall(DMGlobalToLocal(ctx->dm_buoyancy, buoyancy, INSERT_VALUES, buoyancy_local));
468 PetscCall(DMStagVecGetArrayRead(ctx->dm_buoyancy, buoyancy_local, (void *)&arr_buoyancy));
469
470 /* Prepare read-only access to stress data */
471 PetscCall(DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_ELEMENT, 0, &slot_txx));
472 PetscCall(DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_ELEMENT, 1, &slot_tyy));
473 PetscCall(DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_ELEMENT, 2, &slot_tzz));
474 PetscCall(DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_UP_LEFT, 0, &slot_txy_upleft));
475 PetscCall(DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_DOWN_LEFT, 0, &slot_txy_downleft));
476 PetscCall(DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_DOWN_RIGHT, 0, &slot_txy_downright));
477 PetscCall(DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_BACK_LEFT, 0, &slot_txz_backleft));
478 PetscCall(DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_BACK_RIGHT, 0, &slot_txz_backright));
479 PetscCall(DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_FRONT_LEFT, 0, &slot_txz_frontleft));
480 PetscCall(DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_BACK_DOWN, 0, &slot_tyz_backdown));
481 PetscCall(DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_BACK_UP, 0, &slot_tyz_backup));
482 PetscCall(DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_FRONT_DOWN, 0, &slot_tyz_frontdown));
483 PetscCall(DMGetLocalVector(ctx->dm_stress, &stress_local));
484 PetscCall(DMGlobalToLocal(ctx->dm_stress, stress, INSERT_VALUES, stress_local));
485 PetscCall(DMStagVecGetArrayRead(ctx->dm_stress, stress_local, (void *)&arr_stress));
486
487 /* Prepare read-write access to velocity data */
488 PetscCall(DMStagGetLocationSlot(ctx->dm_velocity, DMSTAG_LEFT, 0, &slot_vx_left));
489 PetscCall(DMStagGetLocationSlot(ctx->dm_velocity, DMSTAG_DOWN, 0, &slot_vy_down));
490 PetscCall(DMStagGetLocationSlot(ctx->dm_velocity, DMSTAG_BACK, 0, &slot_vz_back));
491 PetscCall(DMGetLocalVector(ctx->dm_velocity, &velocity_local));
492 PetscCall(DMGlobalToLocal(ctx->dm_velocity, velocity, INSERT_VALUES, velocity_local));
493 PetscCall(DMStagVecGetArray(ctx->dm_velocity, velocity_local, &arr_velocity));
494
495 /* Prepare read-only access to coordinate data */
496 PetscCall(DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_LEFT, &slot_coord_prev));
497 PetscCall(DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_RIGHT, &slot_coord_next));
498 PetscCall(DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_ELEMENT, &slot_coord_element));
499 PetscCall(DMStagGetProductCoordinateArrays(ctx->dm_velocity, (void *)&arr_coord_x, (void *)&arr_coord_y, (void *)&arr_coord_z));
500
501 /* Iterate over interior of the domain, updating the velocities */
502 PetscCall(DMStagGetCorners(ctx->dm_velocity, &startx, &starty, &startz, &nx, &ny, &nz, NULL, NULL, NULL));
503 for (ez = startz; ez < startz + nz; ++ez) {
504 for (ey = starty; ey < starty + ny; ++ey) {
505 for (ex = startx; ex < startx + nx; ++ex) {
506 /* Update y-velocity */
507 if (ey > 0) {
508 const PetscScalar dx = arr_coord_x[ex][slot_coord_next] - arr_coord_x[ex][slot_coord_prev];
509 const PetscScalar dy = arr_coord_y[ey][slot_coord_element] - arr_coord_y[ey - 1][slot_coord_element];
510 const PetscScalar dz = arr_coord_z[ez][slot_coord_next] - arr_coord_z[ez][slot_coord_prev];
511 const PetscScalar B = arr_buoyancy[ez][ey][ex][slot_buoyancy_down];
512
513 arr_velocity[ez][ey][ex][slot_vy_down] += B * ctx->dt * ((arr_stress[ez][ey][ex][slot_txy_downright] - arr_stress[ez][ey][ex][slot_txy_downleft]) / dx + (arr_stress[ez][ey][ex][slot_tyy] - arr_stress[ez][ey - 1][ex][slot_tyy]) / dy + (arr_stress[ez][ey][ex][slot_tyz_frontdown] - arr_stress[ez][ey][ex][slot_tyz_backdown]) / dz);
514 }
515
516 /* Update x-velocity */
517 if (ex > 0) {
518 const PetscScalar dx = arr_coord_x[ex][slot_coord_element] - arr_coord_x[ex - 1][slot_coord_element];
519 const PetscScalar dy = arr_coord_y[ey][slot_coord_next] - arr_coord_y[ey][slot_coord_prev];
520 const PetscScalar dz = arr_coord_z[ez][slot_coord_next] - arr_coord_z[ez][slot_coord_prev];
521 const PetscScalar B = arr_buoyancy[ez][ey][ex][slot_buoyancy_left];
522
523 arr_velocity[ez][ey][ex][slot_vx_left] += B * ctx->dt * ((arr_stress[ez][ey][ex][slot_txx] - arr_stress[ez][ey][ex - 1][slot_txx]) / dx + (arr_stress[ez][ey][ex][slot_txy_upleft] - arr_stress[ez][ey][ex][slot_txy_downleft]) / dy + (arr_stress[ez][ey][ex][slot_txz_frontleft] - arr_stress[ez][ey][ex][slot_txz_backleft]) / dz);
524 }
525
526 /* Update z-velocity */
527 if (ez > 0) {
528 const PetscScalar dx = arr_coord_x[ex][slot_coord_next] - arr_coord_x[ex][slot_coord_prev];
529 const PetscScalar dy = arr_coord_y[ey][slot_coord_next] - arr_coord_y[ey][slot_coord_prev];
530 const PetscScalar dz = arr_coord_z[ez][slot_coord_element] - arr_coord_z[ez - 1][slot_coord_element];
531 const PetscScalar B = arr_buoyancy[ez][ey][ex][slot_buoyancy_back];
532
533 arr_velocity[ez][ey][ex][slot_vz_back] += B * ctx->dt * ((arr_stress[ez][ey][ex][slot_txz_backright] - arr_stress[ez][ey][ex][slot_txz_backleft]) / dx + (arr_stress[ez][ey][ex][slot_tyz_backup] - arr_stress[ez][ey][ex][slot_tyz_backdown]) / dy + (arr_stress[ez][ey][ex][slot_tzz] - arr_stress[ez - 1][ey][ex][slot_tzz]) / dz);
534 }
535 }
536 }
537 }
538
539 /* Restore all access */
540 PetscCall(DMStagVecRestoreArrayRead(ctx->dm_buoyancy, buoyancy_local, (void *)&arr_buoyancy));
541 PetscCall(DMRestoreLocalVector(ctx->dm_buoyancy, &buoyancy_local));
542 PetscCall(DMStagVecRestoreArrayRead(ctx->dm_stress, stress_local, (void *)&arr_stress));
543 PetscCall(DMRestoreLocalVector(ctx->dm_stress, &stress_local));
544 PetscCall(DMStagVecRestoreArray(ctx->dm_velocity, velocity_local, &arr_velocity));
545 PetscCall(DMLocalToGlobal(ctx->dm_velocity, velocity_local, INSERT_VALUES, velocity));
546 PetscCall(DMRestoreLocalVector(ctx->dm_velocity, &velocity_local));
547 PetscCall(DMStagRestoreProductCoordinateArrays(ctx->dm_velocity, (void *)&arr_coord_x, (void *)&arr_coord_y, (void *)&arr_coord_z));
548 PetscFunctionReturn(PETSC_SUCCESS);
549 }
550
UpdateVelocity(const Ctx * ctx,Vec velocity,Vec stress,Vec buoyancy)551 static PetscErrorCode UpdateVelocity(const Ctx *ctx, Vec velocity, Vec stress, Vec buoyancy)
552 {
553 PetscFunctionBeginUser;
554 if (ctx->dim == 2) {
555 PetscCall(UpdateVelocity_2d(ctx, velocity, stress, buoyancy));
556 } else if (ctx->dim == 3) {
557 PetscCall(UpdateVelocity_3d(ctx, velocity, stress, buoyancy));
558 } else SETERRQ(PetscObjectComm((PetscObject)ctx->dm_velocity), PETSC_ERR_SUP, "Unsupported dim %" PetscInt_FMT, ctx->dim);
559 PetscFunctionReturn(PETSC_SUCCESS);
560 }
561
UpdateStress_2d(const Ctx * ctx,Vec velocity,Vec stress,Vec lame)562 static PetscErrorCode UpdateStress_2d(const Ctx *ctx, Vec velocity, Vec stress, Vec lame)
563 {
564 Vec velocity_local, stress_local, lame_local;
565 PetscInt ex, ey, startx, starty, nx, ny;
566 PetscInt slot_coord_next, slot_coord_element, slot_coord_prev;
567 PetscInt slot_vx_left, slot_vy_down, slot_vx_right, slot_vy_up;
568 PetscInt slot_mu_element, slot_lambda_element, slot_mu_downleft;
569 PetscInt slot_txx, slot_tyy, slot_txy_downleft;
570 const PetscScalar **arr_coord_x, **arr_coord_y;
571 const PetscScalar ***arr_velocity, ***arr_lame;
572 PetscScalar ***arr_stress;
573
574 PetscFunctionBeginUser;
575 /* Prepare read-write access to stress data */
576 PetscCall(DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_ELEMENT, 0, &slot_txx));
577 PetscCall(DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_ELEMENT, 1, &slot_tyy));
578 PetscCall(DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_DOWN_LEFT, 0, &slot_txy_downleft));
579 PetscCall(DMGetLocalVector(ctx->dm_stress, &stress_local));
580 PetscCall(DMGlobalToLocal(ctx->dm_stress, stress, INSERT_VALUES, stress_local));
581 PetscCall(DMStagVecGetArray(ctx->dm_stress, stress_local, &arr_stress));
582
583 /* Prepare read-only access to velocity data */
584 PetscCall(DMStagGetLocationSlot(ctx->dm_velocity, DMSTAG_LEFT, 0, &slot_vx_left));
585 PetscCall(DMStagGetLocationSlot(ctx->dm_velocity, DMSTAG_RIGHT, 0, &slot_vx_right));
586 PetscCall(DMStagGetLocationSlot(ctx->dm_velocity, DMSTAG_DOWN, 0, &slot_vy_down));
587 PetscCall(DMStagGetLocationSlot(ctx->dm_velocity, DMSTAG_UP, 0, &slot_vy_up));
588 PetscCall(DMGetLocalVector(ctx->dm_velocity, &velocity_local));
589 PetscCall(DMGlobalToLocal(ctx->dm_velocity, velocity, INSERT_VALUES, velocity_local));
590 PetscCall(DMStagVecGetArrayRead(ctx->dm_velocity, velocity_local, (void *)&arr_velocity));
591
592 /* Prepare read-only access to Lame' data */
593 PetscCall(DMStagGetLocationSlot(ctx->dm_lame, DMSTAG_ELEMENT, 0, &slot_lambda_element));
594 PetscCall(DMStagGetLocationSlot(ctx->dm_lame, DMSTAG_ELEMENT, 1, &slot_mu_element));
595 PetscCall(DMStagGetLocationSlot(ctx->dm_lame, DMSTAG_DOWN_LEFT, 0, &slot_mu_downleft));
596 PetscCall(DMGetLocalVector(ctx->dm_lame, &lame_local));
597 PetscCall(DMGlobalToLocal(ctx->dm_lame, lame, INSERT_VALUES, lame_local));
598 PetscCall(DMStagVecGetArrayRead(ctx->dm_lame, lame_local, (void *)&arr_lame));
599
600 /* Prepare read-only access to coordinate data */
601 PetscCall(DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_LEFT, &slot_coord_prev));
602 PetscCall(DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_RIGHT, &slot_coord_next));
603 PetscCall(DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_ELEMENT, &slot_coord_element));
604 PetscCall(DMStagGetProductCoordinateArrays(ctx->dm_velocity, (void *)&arr_coord_x, (void *)&arr_coord_y, NULL));
605
606 /* Iterate over the interior of the domain, updating the stresses */
607 PetscCall(DMStagGetCorners(ctx->dm_velocity, &startx, &starty, NULL, &nx, &ny, NULL, NULL, NULL, NULL));
608 for (ey = starty; ey < starty + ny; ++ey) {
609 for (ex = startx; ex < startx + nx; ++ex) {
610 /* Update tau_xx and tau_yy*/
611 {
612 const PetscScalar dx = arr_coord_x[ex][slot_coord_next] - arr_coord_x[ex][slot_coord_prev];
613 const PetscScalar dy = arr_coord_y[ey][slot_coord_next] - arr_coord_y[ey][slot_coord_prev];
614 const PetscScalar L = arr_lame[ey][ex][slot_lambda_element];
615 const PetscScalar Lp2M = arr_lame[ey][ex][slot_lambda_element] + 2 * arr_lame[ey][ex][slot_mu_element];
616
617 arr_stress[ey][ex][slot_txx] += Lp2M * ctx->dt * (arr_velocity[ey][ex][slot_vx_right] - arr_velocity[ey][ex][slot_vx_left]) / dx + L * ctx->dt * (arr_velocity[ey][ex][slot_vy_up] - arr_velocity[ey][ex][slot_vy_down]) / dy;
618
619 arr_stress[ey][ex][slot_tyy] += Lp2M * ctx->dt * (arr_velocity[ey][ex][slot_vy_up] - arr_velocity[ey][ex][slot_vy_down]) / dy + L * ctx->dt * (arr_velocity[ey][ex][slot_vx_right] - arr_velocity[ey][ex][slot_vx_left]) / dx;
620 }
621
622 /* Update tau_xy */
623 if (ex > 0 && ey > 0) {
624 const PetscScalar dx = arr_coord_x[ex][slot_coord_element] - arr_coord_x[ex - 1][slot_coord_element];
625 const PetscScalar dy = arr_coord_y[ey][slot_coord_element] - arr_coord_y[ey - 1][slot_coord_element];
626 const PetscScalar M = arr_lame[ey][ex][slot_mu_downleft];
627
628 arr_stress[ey][ex][slot_txy_downleft] += M * ctx->dt * ((arr_velocity[ey][ex][slot_vx_left] - arr_velocity[ey - 1][ex][slot_vx_left]) / dy + (arr_velocity[ey][ex][slot_vy_down] - arr_velocity[ey][ex - 1][slot_vy_down]) / dx);
629 }
630 }
631 }
632
633 /* Restore all access */
634 PetscCall(DMStagVecRestoreArray(ctx->dm_stress, stress_local, &arr_stress));
635 PetscCall(DMLocalToGlobal(ctx->dm_stress, stress_local, INSERT_VALUES, stress));
636 PetscCall(DMRestoreLocalVector(ctx->dm_stress, &stress_local));
637 PetscCall(DMStagVecRestoreArrayRead(ctx->dm_velocity, velocity_local, (void *)&arr_velocity));
638 PetscCall(DMRestoreLocalVector(ctx->dm_velocity, &velocity_local));
639 PetscCall(DMStagVecRestoreArrayRead(ctx->dm_lame, lame_local, (void *)&arr_lame));
640 PetscCall(DMRestoreLocalVector(ctx->dm_lame, &lame_local));
641 PetscCall(DMStagRestoreProductCoordinateArrays(ctx->dm_velocity, (void *)&arr_coord_x, (void *)&arr_coord_y, NULL));
642 PetscFunctionReturn(PETSC_SUCCESS);
643 }
644
UpdateStress_3d(const Ctx * ctx,Vec velocity,Vec stress,Vec lame)645 static PetscErrorCode UpdateStress_3d(const Ctx *ctx, Vec velocity, Vec stress, Vec lame)
646 {
647 Vec velocity_local, stress_local, lame_local;
648 PetscInt ex, ey, ez, startx, starty, startz, nx, ny, nz;
649 PetscInt slot_coord_next, slot_coord_element, slot_coord_prev;
650 PetscInt slot_vx_left, slot_vy_down, slot_vx_right, slot_vy_up, slot_vz_back, slot_vz_front;
651 PetscInt slot_mu_element, slot_lambda_element, slot_mu_downleft, slot_mu_backdown, slot_mu_backleft;
652 PetscInt slot_txx, slot_tyy, slot_tzz, slot_txy_downleft, slot_txz_backleft, slot_tyz_backdown;
653 const PetscScalar **arr_coord_x, **arr_coord_y, **arr_coord_z;
654 const PetscScalar ****arr_velocity, ****arr_lame;
655 PetscScalar ****arr_stress;
656
657 PetscFunctionBeginUser;
658 /* Prepare read-write access to stress data */
659 PetscCall(DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_ELEMENT, 0, &slot_txx));
660 PetscCall(DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_ELEMENT, 1, &slot_tyy));
661 PetscCall(DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_ELEMENT, 2, &slot_tzz));
662 PetscCall(DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_DOWN_LEFT, 0, &slot_txy_downleft));
663 PetscCall(DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_BACK_LEFT, 0, &slot_txz_backleft));
664 PetscCall(DMStagGetLocationSlot(ctx->dm_stress, DMSTAG_BACK_DOWN, 0, &slot_tyz_backdown));
665 PetscCall(DMGetLocalVector(ctx->dm_stress, &stress_local));
666 PetscCall(DMGlobalToLocal(ctx->dm_stress, stress, INSERT_VALUES, stress_local));
667 PetscCall(DMStagVecGetArray(ctx->dm_stress, stress_local, &arr_stress));
668
669 /* Prepare read-only access to velocity data */
670 PetscCall(DMStagGetLocationSlot(ctx->dm_velocity, DMSTAG_LEFT, 0, &slot_vx_left));
671 PetscCall(DMStagGetLocationSlot(ctx->dm_velocity, DMSTAG_RIGHT, 0, &slot_vx_right));
672 PetscCall(DMStagGetLocationSlot(ctx->dm_velocity, DMSTAG_DOWN, 0, &slot_vy_down));
673 PetscCall(DMStagGetLocationSlot(ctx->dm_velocity, DMSTAG_UP, 0, &slot_vy_up));
674 PetscCall(DMStagGetLocationSlot(ctx->dm_velocity, DMSTAG_BACK, 0, &slot_vz_back));
675 PetscCall(DMStagGetLocationSlot(ctx->dm_velocity, DMSTAG_FRONT, 0, &slot_vz_front));
676 PetscCall(DMGetLocalVector(ctx->dm_velocity, &velocity_local));
677 PetscCall(DMGlobalToLocal(ctx->dm_velocity, velocity, INSERT_VALUES, velocity_local));
678 PetscCall(DMStagVecGetArrayRead(ctx->dm_velocity, velocity_local, (void *)&arr_velocity));
679
680 /* Prepare read-only access to Lame' data */
681 PetscCall(DMStagGetLocationSlot(ctx->dm_lame, DMSTAG_ELEMENT, 0, &slot_lambda_element));
682 PetscCall(DMStagGetLocationSlot(ctx->dm_lame, DMSTAG_ELEMENT, 1, &slot_mu_element));
683 PetscCall(DMStagGetLocationSlot(ctx->dm_lame, DMSTAG_DOWN_LEFT, 0, &slot_mu_downleft));
684 PetscCall(DMStagGetLocationSlot(ctx->dm_lame, DMSTAG_BACK_LEFT, 0, &slot_mu_backleft));
685 PetscCall(DMStagGetLocationSlot(ctx->dm_lame, DMSTAG_BACK_DOWN, 0, &slot_mu_backdown));
686 PetscCall(DMGetLocalVector(ctx->dm_lame, &lame_local));
687 PetscCall(DMGlobalToLocal(ctx->dm_lame, lame, INSERT_VALUES, lame_local));
688 PetscCall(DMStagVecGetArrayRead(ctx->dm_lame, lame_local, (void *)&arr_lame));
689
690 /* Prepare read-only access to coordinate data */
691 PetscCall(DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_LEFT, &slot_coord_prev));
692 PetscCall(DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_RIGHT, &slot_coord_next));
693 PetscCall(DMStagGetProductCoordinateLocationSlot(ctx->dm_velocity, DMSTAG_ELEMENT, &slot_coord_element));
694 PetscCall(DMStagGetProductCoordinateArrays(ctx->dm_velocity, (void *)&arr_coord_x, (void *)&arr_coord_y, (void *)&arr_coord_z));
695
696 /* Iterate over the interior of the domain, updating the stresses */
697 PetscCall(DMStagGetCorners(ctx->dm_velocity, &startx, &starty, &startz, &nx, &ny, &nz, NULL, NULL, NULL));
698 for (ez = startz; ez < startz + nz; ++ez) {
699 for (ey = starty; ey < starty + ny; ++ey) {
700 for (ex = startx; ex < startx + nx; ++ex) {
701 /* Update tau_xx, tau_yy, and tau_zz*/
702 {
703 const PetscScalar dx = arr_coord_x[ex][slot_coord_next] - arr_coord_x[ex][slot_coord_prev];
704 const PetscScalar dy = arr_coord_y[ey][slot_coord_next] - arr_coord_y[ey][slot_coord_prev];
705 const PetscScalar dz = arr_coord_z[ez][slot_coord_next] - arr_coord_z[ez][slot_coord_prev];
706 const PetscScalar L = arr_lame[ez][ey][ex][slot_lambda_element];
707 const PetscScalar Lp2M = arr_lame[ez][ey][ex][slot_lambda_element] + 2 * arr_lame[ez][ey][ex][slot_mu_element];
708
709 arr_stress[ez][ey][ex][slot_txx] += Lp2M * ctx->dt * (arr_velocity[ez][ey][ex][slot_vx_right] - arr_velocity[ez][ey][ex][slot_vx_left]) / dx + L * ctx->dt * (arr_velocity[ez][ey][ex][slot_vy_up] - arr_velocity[ez][ey][ex][slot_vy_down]) / dy +
710 L * ctx->dt * (arr_velocity[ez][ey][ex][slot_vz_front] - arr_velocity[ez][ey][ex][slot_vz_back]) / dz;
711
712 arr_stress[ez][ey][ex][slot_tyy] += Lp2M * ctx->dt * (arr_velocity[ez][ey][ex][slot_vy_up] - arr_velocity[ez][ey][ex][slot_vy_down]) / dy + L * ctx->dt * (arr_velocity[ez][ey][ex][slot_vz_front] - arr_velocity[ez][ey][ex][slot_vz_back]) / dz +
713 L * ctx->dt * (arr_velocity[ez][ey][ex][slot_vx_right] - arr_velocity[ez][ey][ex][slot_vx_left]) / dx;
714
715 arr_stress[ez][ey][ex][slot_tzz] += Lp2M * ctx->dt * (arr_velocity[ez][ey][ex][slot_vz_front] - arr_velocity[ez][ey][ex][slot_vz_back]) / dz + L * ctx->dt * (arr_velocity[ez][ey][ex][slot_vx_right] - arr_velocity[ez][ey][ex][slot_vx_left]) / dx +
716 L * ctx->dt * (arr_velocity[ez][ey][ex][slot_vy_up] - arr_velocity[ez][ey][ex][slot_vy_down]) / dy;
717 }
718
719 /* Update tau_xy, tau_xz, tau_yz */
720 {
721 PetscScalar dx = 1.0, dy = 1.0, dz = 1.0; /* initialization to prevent incorrect compiler warnings */
722
723 if (ex > 0) dx = arr_coord_x[ex][slot_coord_element] - arr_coord_x[ex - 1][slot_coord_element];
724 if (ey > 0) dy = arr_coord_y[ey][slot_coord_element] - arr_coord_y[ey - 1][slot_coord_element];
725 if (ez > 0) dz = arr_coord_z[ez][slot_coord_element] - arr_coord_z[ez - 1][slot_coord_element];
726
727 if (ex > 0 && ey > 0) {
728 const PetscScalar M = arr_lame[ez][ey][ex][slot_mu_downleft];
729
730 arr_stress[ez][ey][ex][slot_txy_downleft] += M * ctx->dt * ((arr_velocity[ez][ey][ex][slot_vx_left] - arr_velocity[ez][ey - 1][ex][slot_vx_left]) / dy + (arr_velocity[ez][ey][ex][slot_vy_down] - arr_velocity[ez][ey][ex - 1][slot_vy_down]) / dx);
731 }
732
733 /* Update tau_xz */
734 if (ex > 0 && ez > 0) {
735 const PetscScalar M = arr_lame[ez][ey][ex][slot_mu_backleft];
736
737 arr_stress[ez][ey][ex][slot_txz_backleft] += M * ctx->dt * ((arr_velocity[ez][ey][ex][slot_vx_left] - arr_velocity[ez - 1][ey][ex][slot_vx_left]) / dz + (arr_velocity[ez][ey][ex][slot_vz_back] - arr_velocity[ez][ey][ex - 1][slot_vz_back]) / dx);
738 }
739
740 /* Update tau_yz */
741 if (ey > 0 && ez > 0) {
742 const PetscScalar M = arr_lame[ez][ey][ex][slot_mu_backdown];
743
744 arr_stress[ez][ey][ex][slot_tyz_backdown] += M * ctx->dt * ((arr_velocity[ez][ey][ex][slot_vy_down] - arr_velocity[ez - 1][ey][ex][slot_vy_down]) / dz + (arr_velocity[ez][ey][ex][slot_vz_back] - arr_velocity[ez][ey - 1][ex][slot_vz_back]) / dy);
745 }
746 }
747 }
748 }
749 }
750
751 /* Restore all access */
752 PetscCall(DMStagVecRestoreArray(ctx->dm_stress, stress_local, &arr_stress));
753 PetscCall(DMLocalToGlobal(ctx->dm_stress, stress_local, INSERT_VALUES, stress));
754 PetscCall(DMRestoreLocalVector(ctx->dm_stress, &stress_local));
755 PetscCall(DMStagVecRestoreArrayRead(ctx->dm_velocity, velocity_local, (void *)&arr_velocity));
756 PetscCall(DMRestoreLocalVector(ctx->dm_velocity, &velocity_local));
757 PetscCall(DMStagVecRestoreArrayRead(ctx->dm_lame, lame_local, (void *)&arr_lame));
758 PetscCall(DMRestoreLocalVector(ctx->dm_lame, &lame_local));
759 PetscCall(DMStagRestoreProductCoordinateArrays(ctx->dm_velocity, (void *)&arr_coord_x, (void *)&arr_coord_y, (void *)&arr_coord_z));
760 PetscFunctionReturn(PETSC_SUCCESS);
761 }
762
UpdateStress(const Ctx * ctx,Vec velocity,Vec stress,Vec lame)763 static PetscErrorCode UpdateStress(const Ctx *ctx, Vec velocity, Vec stress, Vec lame)
764 {
765 PetscFunctionBeginUser;
766 if (ctx->dim == 2) {
767 PetscCall(UpdateStress_2d(ctx, velocity, stress, lame));
768 } else if (ctx->dim == 3) {
769 PetscCall(UpdateStress_3d(ctx, velocity, stress, lame));
770 }
771 PetscFunctionReturn(PETSC_SUCCESS);
772 }
773
DumpStress(const Ctx * ctx,Vec stress,PetscInt timestep)774 static PetscErrorCode DumpStress(const Ctx *ctx, Vec stress, PetscInt timestep)
775 {
776 DM da_normal, da_shear = NULL;
777 Vec vec_normal, vec_shear = NULL;
778
779 PetscFunctionBeginUser;
780 PetscCall(DMStagVecSplitToDMDA(ctx->dm_stress, stress, DMSTAG_ELEMENT, -ctx->dim, &da_normal, &vec_normal));
781 PetscCall(PetscObjectSetName((PetscObject)vec_normal, "normal stresses"));
782
783 /* Dump element-based fields to a .vtr file */
784 {
785 PetscViewer viewer;
786 char filename[PETSC_MAX_PATH_LEN];
787
788 PetscCall(PetscSNPrintf(filename, sizeof(filename), "ex6_stress_normal_%.4" PetscInt_FMT ".vtr", timestep));
789 PetscCall(PetscViewerVTKOpen(PetscObjectComm((PetscObject)da_normal), filename, FILE_MODE_WRITE, &viewer));
790 PetscCall(VecView(vec_normal, viewer));
791 PetscCall(PetscViewerDestroy(&viewer));
792 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Created %s\n", filename));
793 }
794
795 if (ctx->dim == 2) {
796 /* (2D only) Dump vertex-based fields to a second .vtr file */
797 PetscCall(DMStagVecSplitToDMDA(ctx->dm_stress, stress, DMSTAG_DOWN_LEFT, 0, &da_shear, &vec_shear));
798 PetscCall(PetscObjectSetName((PetscObject)vec_shear, "shear stresses"));
799
800 {
801 PetscViewer viewer;
802 char filename[PETSC_MAX_PATH_LEN];
803
804 PetscCall(PetscSNPrintf(filename, sizeof(filename), "ex6_stress_shear_%.4" PetscInt_FMT ".vtr", timestep));
805 PetscCall(PetscViewerVTKOpen(PetscObjectComm((PetscObject)da_normal), filename, FILE_MODE_WRITE, &viewer));
806 PetscCall(VecView(vec_shear, viewer));
807 PetscCall(PetscViewerDestroy(&viewer));
808 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Created %s\n", filename));
809 }
810 }
811
812 /* Destroy DMDAs and Vecs */
813 PetscCall(DMDestroy(&da_normal));
814 PetscCall(DMDestroy(&da_shear));
815 PetscCall(VecDestroy(&vec_normal));
816 PetscCall(VecDestroy(&vec_shear));
817 PetscFunctionReturn(PETSC_SUCCESS);
818 }
819
DumpVelocity(const Ctx * ctx,Vec velocity,PetscInt timestep)820 static PetscErrorCode DumpVelocity(const Ctx *ctx, Vec velocity, PetscInt timestep)
821 {
822 DM dmVelAvg;
823 Vec velAvg;
824 DM daVelAvg;
825 Vec vecVelAvg;
826 Vec velocity_local;
827 PetscInt ex, ey, ez, startx, starty, startz, nx, ny, nz;
828
829 PetscFunctionBeginUser;
830 if (ctx->dim == 2) {
831 PetscCall(DMStagCreateCompatibleDMStag(ctx->dm_velocity, 0, 0, 2, 0, &dmVelAvg)); /* 2 dof per element */
832 } else if (ctx->dim == 3) {
833 PetscCall(DMStagCreateCompatibleDMStag(ctx->dm_velocity, 0, 0, 0, 3, &dmVelAvg)); /* 3 dof per element */
834 } else SETERRQ(PetscObjectComm((PetscObject)ctx->dm_velocity), PETSC_ERR_SUP, "Unsupported dim %" PetscInt_FMT, ctx->dim);
835 PetscCall(DMSetUp(dmVelAvg));
836 PetscCall(DMStagSetUniformCoordinatesProduct(dmVelAvg, ctx->xmin, ctx->xmax, ctx->ymin, ctx->ymax, ctx->zmin, ctx->zmax));
837 PetscCall(DMCreateGlobalVector(dmVelAvg, &velAvg));
838 PetscCall(DMGetLocalVector(ctx->dm_velocity, &velocity_local));
839 PetscCall(DMGlobalToLocal(ctx->dm_velocity, velocity, INSERT_VALUES, velocity_local));
840 PetscCall(DMStagGetCorners(dmVelAvg, &startx, &starty, &startz, &nx, &ny, &nz, NULL, NULL, NULL));
841 if (ctx->dim == 2) {
842 for (ey = starty; ey < starty + ny; ++ey) {
843 for (ex = startx; ex < startx + nx; ++ex) {
844 DMStagStencil from[4], to[2];
845 PetscScalar valFrom[4], valTo[2];
846
847 from[0].i = ex;
848 from[0].j = ey;
849 from[0].loc = DMSTAG_UP;
850 from[0].c = 0;
851 from[1].i = ex;
852 from[1].j = ey;
853 from[1].loc = DMSTAG_DOWN;
854 from[1].c = 0;
855 from[2].i = ex;
856 from[2].j = ey;
857 from[2].loc = DMSTAG_LEFT;
858 from[2].c = 0;
859 from[3].i = ex;
860 from[3].j = ey;
861 from[3].loc = DMSTAG_RIGHT;
862 from[3].c = 0;
863 PetscCall(DMStagVecGetValuesStencil(ctx->dm_velocity, velocity_local, 4, from, valFrom));
864 to[0].i = ex;
865 to[0].j = ey;
866 to[0].loc = DMSTAG_ELEMENT;
867 to[0].c = 0;
868 valTo[0] = 0.5 * (valFrom[2] + valFrom[3]);
869 to[1].i = ex;
870 to[1].j = ey;
871 to[1].loc = DMSTAG_ELEMENT;
872 to[1].c = 1;
873 valTo[1] = 0.5 * (valFrom[0] + valFrom[1]);
874 PetscCall(DMStagVecSetValuesStencil(dmVelAvg, velAvg, 2, to, valTo, INSERT_VALUES));
875 }
876 }
877 } else if (ctx->dim == 3) {
878 for (ez = startz; ez < startz + nz; ++ez) {
879 for (ey = starty; ey < starty + ny; ++ey) {
880 for (ex = startx; ex < startx + nx; ++ex) {
881 DMStagStencil from[6], to[3];
882 PetscScalar valFrom[6], valTo[3];
883
884 from[0].i = ex;
885 from[0].j = ey;
886 from[0].k = ez;
887 from[0].loc = DMSTAG_UP;
888 from[0].c = 0;
889 from[1].i = ex;
890 from[1].j = ey;
891 from[1].k = ez;
892 from[1].loc = DMSTAG_DOWN;
893 from[1].c = 0;
894 from[2].i = ex;
895 from[2].j = ey;
896 from[2].k = ez;
897 from[2].loc = DMSTAG_LEFT;
898 from[2].c = 0;
899 from[3].i = ex;
900 from[3].j = ey;
901 from[3].k = ez;
902 from[3].loc = DMSTAG_RIGHT;
903 from[3].c = 0;
904 from[4].i = ex;
905 from[4].j = ey;
906 from[4].k = ez;
907 from[4].loc = DMSTAG_FRONT;
908 from[4].c = 0;
909 from[5].i = ex;
910 from[5].j = ey;
911 from[5].k = ez;
912 from[5].loc = DMSTAG_BACK;
913 from[5].c = 0;
914 PetscCall(DMStagVecGetValuesStencil(ctx->dm_velocity, velocity_local, 6, from, valFrom));
915 to[0].i = ex;
916 to[0].j = ey;
917 to[0].k = ez;
918 to[0].loc = DMSTAG_ELEMENT;
919 to[0].c = 0;
920 valTo[0] = 0.5 * (valFrom[2] + valFrom[3]);
921 to[1].i = ex;
922 to[1].j = ey;
923 to[1].k = ez;
924 to[1].loc = DMSTAG_ELEMENT;
925 to[1].c = 1;
926 valTo[1] = 0.5 * (valFrom[0] + valFrom[1]);
927 to[2].i = ex;
928 to[2].j = ey;
929 to[2].k = ez;
930 to[2].loc = DMSTAG_ELEMENT;
931 to[2].c = 2;
932 valTo[2] = 0.5 * (valFrom[4] + valFrom[5]);
933 PetscCall(DMStagVecSetValuesStencil(dmVelAvg, velAvg, 3, to, valTo, INSERT_VALUES));
934 }
935 }
936 }
937 } else SETERRQ(PetscObjectComm((PetscObject)ctx->dm_velocity), PETSC_ERR_SUP, "Unsupported dim %" PetscInt_FMT, ctx->dim);
938 PetscCall(VecAssemblyBegin(velAvg));
939 PetscCall(VecAssemblyEnd(velAvg));
940 PetscCall(DMRestoreLocalVector(ctx->dm_velocity, &velocity_local));
941
942 PetscCall(DMStagVecSplitToDMDA(dmVelAvg, velAvg, DMSTAG_ELEMENT, -3, &daVelAvg, &vecVelAvg)); /* note -3 : pad with zero in 2D case */
943 PetscCall(PetscObjectSetName((PetscObject)vecVelAvg, "Velocity (Averaged)"));
944
945 /* Dump element-based fields to a .vtr file */
946 {
947 PetscViewer viewer;
948 char filename[PETSC_MAX_PATH_LEN];
949
950 PetscCall(PetscSNPrintf(filename, sizeof(filename), "ex6_velavg_%.4" PetscInt_FMT ".vtr", timestep));
951 PetscCall(PetscViewerVTKOpen(PetscObjectComm((PetscObject)daVelAvg), filename, FILE_MODE_WRITE, &viewer));
952 PetscCall(VecView(vecVelAvg, viewer));
953 PetscCall(PetscViewerDestroy(&viewer));
954 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Created %s\n", filename));
955 }
956
957 /* Destroy DMDAs and Vecs */
958 PetscCall(VecDestroy(&vecVelAvg));
959 PetscCall(DMDestroy(&daVelAvg));
960 PetscCall(VecDestroy(&velAvg));
961 PetscCall(DMDestroy(&dmVelAvg));
962 PetscFunctionReturn(PETSC_SUCCESS);
963 }
964
965 /*TEST
966
967 test:
968 suffix: 1
969 requires: !complex
970 nsize: 1
971 args: -stag_grid_x 12 -stag_grid_y 7 -nsteps 4 -dump_output 0
972
973 test:
974 suffix: 2
975 requires: !complex
976 nsize: 9
977 args: -stag_grid_x 12 -stag_grid_y 15 -nsteps 3 -dump_output 0
978
979 test:
980 suffix: 3
981 requires: !complex
982 nsize: 1
983 args: -stag_grid_x 12 -stag_grid_y 7 -stag_grid_z 11 -nsteps 2 -dim 3 -dump_output 0
984
985 test:
986 suffix: 4
987 requires: !complex
988 nsize: 12
989 args: -stag_grid_x 12 -stag_grid_y 15 -stag_grid_z 8 -nsteps 3 -dim 3 -dump_output 0
990
991 TEST*/
992