Actual source code: ex4.c

  1: static char help[] = "Solves the incompressible, variable-viscosity Stokes equation in 2D or 3D, driven by buoyancy variations.\n"
  2:                      "-dim: set dimension (2 or 3)\n"
  3:                      "-nondimensional: replace dimensional domain and coefficients with nondimensional ones\n"
  4:                      "-isoviscous: use constant viscosity\n"
  5:                      "-levels: number of grids to create, by coarsening\n"
  6:                      "-rediscretize: create operators for all grids and set up a FieldSplit/MG solver\n"
  7:                      "-dump_solution: dump VTK files\n\n";

  9: #include <petscdm.h>
 10: #include <petscksp.h>
 11: #include <petscdmstag.h>
 12: #include <petscdmda.h>

 14: /* Application context - grid-based data*/
 15: typedef struct LevelCtxData_ {
 16:   DM          dm_stokes, dm_coefficients, dm_faces;
 17:   Vec         coeff;
 18:   PetscInt    cells_x, cells_y, cells_z; /* redundant with DMs */
 19:   PetscReal   hx_characteristic, hy_characteristic, hz_characteristic;
 20:   PetscScalar K_bound, K_cont;
 21: } LevelCtxData;
 22: typedef LevelCtxData *LevelCtx;

 24: /* Application context - problem and grid(s) (but not solver-specific data) */
 25: typedef struct CtxData_ {
 26:   MPI_Comm    comm;
 27:   PetscInt    dim;                       /* redundant with DMs */
 28:   PetscInt    cells_x, cells_y, cells_z; /* Redundant with finest DMs */
 29:   PetscReal   xmax, ymax, xmin, ymin, zmin, zmax;
 30:   PetscScalar eta1, eta2, rho1, rho2, gy, eta_characteristic;
 31:   PetscBool   pin_pressure;
 32:   PetscScalar (*GetEta)(struct CtxData_ *, PetscScalar, PetscScalar, PetscScalar);
 33:   PetscScalar (*GetRho)(struct CtxData_ *, PetscScalar, PetscScalar, PetscScalar);
 34:   PetscInt  n_levels;
 35:   LevelCtx *levels;
 36: } CtxData;
 37: typedef CtxData *Ctx;

 39: /* Helper to pass system-creation parameters */
 40: typedef struct SystemParameters_ {
 41:   Ctx       ctx;
 42:   PetscInt  level;
 43:   PetscBool include_inverse_visc, faces_only;
 44: } SystemParametersData;
 45: typedef SystemParametersData *SystemParameters;

 47: /* Main logic */
 48: static PetscErrorCode AttachNullspace(DM, Mat);
 49: static PetscErrorCode CreateAuxiliaryOperator(Ctx, PetscInt, Mat *);
 50: static PetscErrorCode CreateSystem(SystemParameters, Mat *, Vec *);
 51: static PetscErrorCode CtxCreateAndSetFromOptions(Ctx *);
 52: static PetscErrorCode CtxDestroy(Ctx *);
 53: static PetscErrorCode DumpSolution(Ctx, PetscInt, Vec);
 54: static PetscErrorCode OperatorInsertInverseViscosityPressureTerms(DM, DM, Vec, PetscScalar, Mat);
 55: static PetscErrorCode PopulateCoefficientData(Ctx, PetscInt);
 56: static PetscErrorCode SystemParametersCreate(SystemParameters *, Ctx);
 57: static PetscErrorCode SystemParametersDestroy(SystemParameters *);

 59: int main(int argc, char **argv)
 60: {
 61:   Ctx       ctx;
 62:   Mat       A, *A_faces = NULL, S_hat = NULL, P = NULL;
 63:   Vec       x, b;
 64:   KSP       ksp;
 65:   DM        dm_stokes, dm_coefficients;
 66:   PetscBool dump_solution, build_auxiliary_operator, rediscretize, custom_pc_mat;

 69:   PetscInitialize(&argc, &argv, (char *)0, help);

 71:   /* Accept options for program behavior */
 72:   dump_solution = PETSC_FALSE;
 73:   PetscOptionsGetBool(NULL, NULL, "-dump_solution", &dump_solution, NULL);
 74:   rediscretize = PETSC_FALSE;
 75:   PetscOptionsGetBool(NULL, NULL, "-rediscretize", &rediscretize, NULL);
 76:   build_auxiliary_operator = rediscretize;
 77:   PetscOptionsGetBool(NULL, NULL, "-build_auxiliary_operator", &build_auxiliary_operator, NULL);
 78:   custom_pc_mat = PETSC_FALSE;
 79:   PetscOptionsGetBool(NULL, NULL, "-custom_pc_mat", &custom_pc_mat, NULL);

 81:   /* Populate application context */
 82:   CtxCreateAndSetFromOptions(&ctx);

 84:   /* Create two DMStag objects, corresponding to the same domain and parallel
 85:      decomposition ("topology"). Each defines a different set of fields on
 86:      the domain ("section"); the first the solution to the Stokes equations
 87:      (x- and y-velocities and scalar pressure), and the second holds coefficients
 88:      (viscosities on elements and viscosities+densities on corners/edges in 2d/3d) */
 89:   if (ctx->dim == 2) {
 90:     PetscCall(DMStagCreate2d(ctx->comm, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, ctx->cells_x, ctx->cells_y, /* Global element counts */
 91:                              PETSC_DECIDE, PETSC_DECIDE,                                                /* Determine parallel decomposition automatically */
 92:                              0, 1, 1,                                                                   /* dof: 0 per vertex, 1 per edge, 1 per face/element */
 93:                              DMSTAG_STENCIL_BOX, 1,                                                     /* elementwise stencil width */
 94:                              NULL, NULL, &ctx->levels[ctx->n_levels - 1]->dm_stokes));
 95:   } else if (ctx->dim == 3) {
 96:     PetscCall(DMStagCreate3d(ctx->comm, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, ctx->cells_x, ctx->cells_y, ctx->cells_z, PETSC_DECIDE, PETSC_DECIDE, PETSC_DECIDE, 0, 0, 1, 1, /* dof: 1 per face, 1 per element */
 97:                              DMSTAG_STENCIL_BOX, 1, NULL, NULL, NULL, &ctx->levels[ctx->n_levels - 1]->dm_stokes));
 98:   } else SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_SUP, "Unsupported dimension: %" PetscInt_FMT, ctx->dim);
 99:   dm_stokes = ctx->levels[ctx->n_levels - 1]->dm_stokes;
100:   DMSetFromOptions(dm_stokes);
101:   DMStagGetGlobalSizes(dm_stokes, &ctx->cells_x, &ctx->cells_y, &ctx->cells_z);
102:   DMSetUp(dm_stokes);
103:   DMStagSetUniformCoordinatesProduct(dm_stokes, ctx->xmin, ctx->xmax, ctx->ymin, ctx->ymax, ctx->zmin, ctx->zmax);

105:   if (ctx->dim == 2) DMStagCreateCompatibleDMStag(dm_stokes, 2, 0, 1, 0, &ctx->levels[ctx->n_levels - 1]->dm_coefficients);
106:   else if (ctx->dim == 3) DMStagCreateCompatibleDMStag(dm_stokes, 0, 2, 0, 1, &ctx->levels[ctx->n_levels - 1]->dm_coefficients);
107:   else SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_SUP, "Unsupported dimension: %" PetscInt_FMT, ctx->dim);
108:   dm_coefficients = ctx->levels[ctx->n_levels - 1]->dm_coefficients;
109:   DMSetUp(dm_coefficients);
110:   DMStagSetUniformCoordinatesProduct(dm_coefficients, ctx->xmin, ctx->xmax, ctx->ymin, ctx->ymax, ctx->zmin, ctx->zmax);

112:   /* Create additional DMs by coarsening. 0 is the coarsest level */
113:   for (PetscInt level = ctx->n_levels - 1; level > 0; --level) {
114:     DMCoarsen(ctx->levels[level]->dm_stokes, MPI_COMM_NULL, &ctx->levels[level - 1]->dm_stokes);
115:     DMCoarsen(ctx->levels[level]->dm_coefficients, MPI_COMM_NULL, &ctx->levels[level - 1]->dm_coefficients);
116:   }

118:   /* Compute scaling constants, knowing grid spacing */
119:   ctx->eta_characteristic = PetscMin(PetscRealPart(ctx->eta1), PetscRealPart(ctx->eta2));
120:   for (PetscInt level = 0; level < ctx->n_levels; ++level) {
121:     PetscInt  N[3];
122:     PetscReal hx_avg_inv;

124:     DMStagGetGlobalSizes(ctx->levels[level]->dm_stokes, &N[0], &N[1], &N[2]);
125:     ctx->levels[level]->hx_characteristic = (ctx->xmax - ctx->xmin) / N[0];
126:     ctx->levels[level]->hy_characteristic = (ctx->ymax - ctx->ymin) / N[1];
127:     ctx->levels[level]->hz_characteristic = (ctx->zmax - ctx->zmin) / N[2];
128:     if (ctx->dim == 2) {
129:       hx_avg_inv = 2.0 / (ctx->levels[level]->hx_characteristic + ctx->levels[level]->hy_characteristic);
130:     } else if (ctx->dim == 3) {
131:       hx_avg_inv = 3.0 / (ctx->levels[level]->hx_characteristic + ctx->levels[level]->hy_characteristic + ctx->levels[level]->hz_characteristic);
132:     } else SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_SUP, "Not Implemented for dimension %" PetscInt_FMT, ctx->dim);
133:     ctx->levels[level]->K_cont  = ctx->eta_characteristic * hx_avg_inv;
134:     ctx->levels[level]->K_bound = ctx->eta_characteristic * hx_avg_inv * hx_avg_inv;
135:   }

137:   /* Populate coefficient data */
138:   for (PetscInt level = 0; level < ctx->n_levels; ++level) PopulateCoefficientData(ctx, level);

140:   /* Construct main system */
141:   {
142:     SystemParameters system_parameters;

144:     SystemParametersCreate(&system_parameters, ctx);
145:     CreateSystem(system_parameters, &A, &b);
146:     SystemParametersDestroy(&system_parameters);
147:   }

149:   /* Attach a constant-pressure nullspace to the fine-level operator */
150:   if (!ctx->pin_pressure) AttachNullspace(dm_stokes, A);

152:   /* Set up solver */
153:   KSPCreate(ctx->comm, &ksp);
154:   KSPSetType(ksp, KSPFGMRES);
155:   {
156:     /* Default to a direct solver, if a package is available */
157:     PetscMPIInt size;

159:     MPI_Comm_size(ctx->comm, &size);
160:     if (PetscDefined(HAVE_SUITESPARSE) && size == 1) {
161:       PC pc;

163:       KSPGetPC(ksp, &pc);
164:       PCSetType(pc, PCLU);
165:       PCFactorSetMatSolverType(pc, MATSOLVERUMFPACK); /* A default, requires SuiteSparse */
166:     }
167:     if (PetscDefined(HAVE_MUMPS) && size > 1) {
168:       PC pc;

170:       KSPGetPC(ksp, &pc);
171:       PCSetType(pc, PCLU);
172:       PCFactorSetMatSolverType(pc, MATSOLVERMUMPS); /* A default, requires MUMPS */
173:     }
174:   }

176:   /* Create and set a custom preconditioning matrix */
177:   if (custom_pc_mat) {
178:     SystemParameters system_parameters;

180:     SystemParametersCreate(&system_parameters, ctx);
181:     system_parameters->include_inverse_visc = PETSC_TRUE;
182:     CreateSystem(system_parameters, &P, NULL);
183:     SystemParametersDestroy(&system_parameters);
184:     KSPSetOperators(ksp, A, P);
185:   } else {
186:     KSPSetOperators(ksp, A, A);
187:   }

189:   KSPSetDM(ksp, dm_stokes);
190:   KSPSetDMActive(ksp, PETSC_FALSE);

192:   /* Finish setting up solver (can override options set above) */
193:   KSPSetFromOptions(ksp);

195:   /* Additional solver configuration that can involve setting up and CANNOT be
196:      overridden from the command line */

198:   /* Construct an auxiliary operator for use a Schur complement preconditioner,
199:      and tell PCFieldSplit to use it (which has no effect if not using that PC) */
200:   if (build_auxiliary_operator && !rediscretize) {
201:     PC pc;

203:     KSPGetPC(ksp, &pc);
204:     CreateAuxiliaryOperator(ctx, ctx->n_levels - 1, &S_hat);
205:     PCFieldSplitSetSchurPre(pc, PC_FIELDSPLIT_SCHUR_PRE_USER, S_hat);
206:     KSPSetFromOptions(ksp);
207:   }

209:   if (rediscretize) {
210:     /* Set up an ABF solver with rediscretized geometric multigrid on the velocity-velocity block */
211:     PC  pc, pc_faces;
212:     KSP ksp_faces;

214:     if (ctx->n_levels < 2) PetscPrintf(ctx->comm, "Warning: not using multiple levels!\n");

216:     KSPGetPC(ksp, &pc);
217:     PCSetType(pc, PCFIELDSPLIT);
218:     PCFieldSplitSetType(pc, PC_COMPOSITE_SCHUR);
219:     PCFieldSplitSetSchurFactType(pc, PC_FIELDSPLIT_SCHUR_FACT_UPPER);
220:     if (build_auxiliary_operator) {
221:       CreateAuxiliaryOperator(ctx, ctx->n_levels - 1, &S_hat);
222:       PCFieldSplitSetSchurPre(pc, PC_FIELDSPLIT_SCHUR_PRE_USER, S_hat);
223:     }

225:     /* Create rediscretized velocity-only DMs and operators */
226:     PetscMalloc1(ctx->n_levels, &A_faces);
227:     for (PetscInt level = 0; level < ctx->n_levels; ++level) {
228:       if (ctx->dim == 2) {
229:         DMStagCreateCompatibleDMStag(ctx->levels[level]->dm_stokes, 0, 1, 0, 0, &ctx->levels[level]->dm_faces);
230:       } else if (ctx->dim == 3) {
231:         DMStagCreateCompatibleDMStag(ctx->levels[level]->dm_stokes, 0, 0, 1, 0, &ctx->levels[level]->dm_faces);
232:       } else SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_SUP, "Not Implemented for dimension %" PetscInt_FMT, ctx->dim);
233:       {
234:         SystemParameters system_parameters;

236:         SystemParametersCreate(&system_parameters, ctx);
237:         system_parameters->faces_only = PETSC_TRUE;
238:         system_parameters->level      = level;
239:         CreateSystem(system_parameters, &A_faces[level], NULL);
240:         SystemParametersDestroy(&system_parameters);
241:       }
242:     }

244:     /* Set up to populate enough to define the sub-solver */
245:     KSPSetUp(ksp);

247:     /* Set multigrid components and settings */
248:     {
249:       KSP *sub_ksp;

251:       PCFieldSplitSchurGetSubKSP(pc, NULL, &sub_ksp);
252:       ksp_faces = sub_ksp[0];
253:       PetscFree(sub_ksp);
254:     }
255:     KSPSetType(ksp_faces, KSPGCR);
256:     KSPGetPC(ksp_faces, &pc_faces);
257:     PCSetType(pc_faces, PCMG);
258:     PCMGSetLevels(pc_faces, ctx->n_levels, NULL);
259:     for (PetscInt level = 0; level < ctx->n_levels; ++level) {
260:       KSP ksp_level;
261:       PC  pc_level;

263:       /* Smoothers */
264:       PCMGGetSmoother(pc_faces, level, &ksp_level);
265:       KSPGetPC(ksp_level, &pc_level);
266:       KSPSetOperators(ksp_level, A_faces[level], A_faces[level]);
267:       if (level > 0) PCSetType(pc_level, PCJACOBI);

269:       /* Transfer Operators */
270:       if (level > 0) {
271:         Mat restriction, interpolation;
272:         DM  dm_level   = ctx->levels[level]->dm_faces;
273:         DM  dm_coarser = ctx->levels[level - 1]->dm_faces;

275:         DMCreateInterpolation(dm_coarser, dm_level, &interpolation, NULL);
276:         PCMGSetInterpolation(pc_faces, level, interpolation);
277:         MatDestroy(&interpolation);
278:         DMCreateRestriction(dm_coarser, dm_level, &restriction);
279:         PCMGSetRestriction(pc_faces, level, restriction);
280:         MatDestroy(&restriction);
281:       }
282:     }
283:   }

285:   /* Solve */
286:   VecDuplicate(b, &x);
287:   KSPSolve(ksp, b, x);
288:   {
289:     KSPConvergedReason reason;
290:     KSPGetConvergedReason(ksp, &reason);
292:   }

294:   /* Dump solution by converting to DMDAs and dumping */
295:   if (dump_solution) DumpSolution(ctx, ctx->n_levels - 1, x);

297:   /* Destroy PETSc objects and finalize */
298:   MatDestroy(&A);
299:   PetscFree(A);
300:   if (A_faces) {
301:     for (PetscInt level = 0; level < ctx->n_levels; ++level) {
302:       if (A_faces[level]) MatDestroy(&A_faces[level]);
303:     }
304:     PetscFree(A_faces);
305:   }
306:   if (P) MatDestroy(&P);
307:   VecDestroy(&x);
308:   VecDestroy(&b);
309:   MatDestroy(&S_hat);
310:   KSPDestroy(&ksp);
311:   CtxDestroy(&ctx);
312:   PetscFinalize();
313:   return 0;
314: }

316: static PetscScalar GetEta_constant(Ctx ctx, PetscScalar x, PetscScalar y, PetscScalar z)
317: {
318:   (void)ctx;
319:   (void)x;
320:   (void)y;
321:   (void)z;
322:   return 1.0;
323: }

325: static PetscScalar GetRho_layers(Ctx ctx, PetscScalar x, PetscScalar y, PetscScalar z)
326: {
327:   (void)y;
328:   (void)z;
329:   return PetscRealPart(x) < (ctx->xmax - ctx->xmin) / 2.0 ? ctx->rho1 : ctx->rho2;
330: }

332: static PetscScalar GetEta_layers(Ctx ctx, PetscScalar x, PetscScalar y, PetscScalar z)
333: {
334:   (void)y;
335:   (void)z;
336:   return PetscRealPart(x) < (ctx->xmax - ctx->xmin) / 2.0 ? ctx->eta1 : ctx->eta2;
337: }

339: static PetscScalar GetRho_sinker_box2(Ctx ctx, PetscScalar x, PetscScalar y, PetscScalar z)
340: {
341:   (void)z;
342:   const PetscReal d  = ctx->xmax - ctx->xmin;
343:   const PetscReal xx = PetscRealPart(x) / d - 0.5;
344:   const PetscReal yy = PetscRealPart(y) / d - 0.5;
345:   return (xx * xx > 0.15 * 0.15 || yy * yy > 0.15 * 0.15) ? ctx->rho1 : ctx->rho2;
346: }

348: static PetscScalar GetEta_sinker_box2(Ctx ctx, PetscScalar x, PetscScalar y, PetscScalar z)
349: {
350:   (void)z;
351:   const PetscReal d  = ctx->xmax - ctx->xmin;
352:   const PetscReal xx = PetscRealPart(x) / d - 0.5;
353:   const PetscReal yy = PetscRealPart(y) / d - 0.5;
354:   return (xx * xx > 0.15 * 0.15 || yy * yy > 0.15 * 0.15) ? ctx->eta1 : ctx->eta2;
355: }

357: static PetscScalar GetRho_sinker_box3(Ctx ctx, PetscScalar x, PetscScalar y, PetscScalar z)
358: {
359:   const PetscReal d          = ctx->xmax - ctx->xmin;
360:   const PetscReal xx         = PetscRealPart(x) / d - 0.5;
361:   const PetscReal yy         = PetscRealPart(y) / d - 0.5;
362:   const PetscReal zz         = PetscRealPart(z) / d - 0.5;
363:   const PetscReal half_width = 0.15;
364:   return (PetscAbsReal(xx) > half_width || PetscAbsReal(yy) > half_width || PetscAbsReal(zz) > half_width) ? ctx->rho1 : ctx->rho2;
365: }

367: static PetscScalar GetEta_sinker_box3(Ctx ctx, PetscScalar x, PetscScalar y, PetscScalar z)
368: {
369:   const PetscReal d          = ctx->xmax - ctx->xmin;
370:   const PetscReal xx         = PetscRealPart(x) / d - 0.5;
371:   const PetscReal yy         = PetscRealPart(y) / d - 0.5;
372:   const PetscReal zz         = PetscRealPart(z) / d - 0.5;
373:   const PetscReal half_width = 0.15;
374:   return (PetscAbsReal(xx) > half_width || PetscAbsReal(yy) > half_width || PetscAbsReal(zz) > half_width) ? ctx->eta1 : ctx->eta2;
375: }

377: static PetscScalar GetRho_sinker_sphere3(Ctx ctx, PetscScalar x, PetscScalar y, PetscScalar z)
378: {
379:   const PetscReal d          = ctx->xmax - ctx->xmin;
380:   const PetscReal xx         = PetscRealPart(x) / d - 0.5;
381:   const PetscReal yy         = PetscRealPart(y) / d - 0.5;
382:   const PetscReal zz         = PetscRealPart(z) / d - 0.5;
383:   const PetscReal half_width = 0.3;
384:   return (xx * xx + yy * yy + zz * zz > half_width * half_width) ? ctx->rho1 : ctx->rho2;
385: }

387: static PetscScalar GetEta_sinker_sphere3(Ctx ctx, PetscScalar x, PetscScalar y, PetscScalar z)
388: {
389:   const PetscReal d          = ctx->xmax - ctx->xmin;
390:   const PetscReal xx         = PetscRealPart(x) / d - 0.5;
391:   const PetscReal yy         = PetscRealPart(y) / d - 0.5;
392:   const PetscReal zz         = PetscRealPart(z) / d - 0.5;
393:   const PetscReal half_width = 0.3;
394:   return (xx * xx + yy * yy + zz * zz > half_width * half_width) ? ctx->eta1 : ctx->eta2;
395: }

397: static PetscScalar GetEta_blob3(Ctx ctx, PetscScalar x, PetscScalar y, PetscScalar z)
398: {
399:   const PetscReal d  = ctx->xmax - ctx->xmin;
400:   const PetscReal xx = PetscRealPart(x) / d - 0.5;
401:   const PetscReal yy = PetscRealPart(y) / d - 0.5;
402:   const PetscReal zz = PetscRealPart(z) / d - 0.5;
403:   return ctx->eta1 + ctx->eta2 * PetscExpScalar(-20.0 * (xx * xx + yy * yy + zz * zz));
404: }

406: static PetscScalar GetRho_blob3(Ctx ctx, PetscScalar x, PetscScalar y, PetscScalar z)
407: {
408:   const PetscReal d  = ctx->xmax - ctx->xmin;
409:   const PetscReal xx = PetscRealPart(x) / d - 0.5;
410:   const PetscReal yy = PetscRealPart(y) / d - 0.5;
411:   const PetscReal zz = PetscRealPart(z) / d - 0.5;
412:   return ctx->rho1 + ctx->rho2 * PetscExpScalar(-20.0 * (xx * xx + yy * yy + zz * zz));
413: }

415: static PetscErrorCode LevelCtxCreate(LevelCtx *p_level_ctx)
416: {
417:   LevelCtx level_ctx;

420:   PetscMalloc1(1, p_level_ctx);
421:   level_ctx                  = *p_level_ctx;
422:   level_ctx->dm_stokes       = NULL;
423:   level_ctx->dm_coefficients = NULL;
424:   level_ctx->dm_faces        = NULL;
425:   return 0;
426: }

428: static PetscErrorCode LevelCtxDestroy(LevelCtx *p_level_ctx)
429: {
430:   LevelCtx level_ctx;

433:   level_ctx = *p_level_ctx;
434:   if (level_ctx->dm_stokes) DMDestroy(&level_ctx->dm_stokes);
435:   if (level_ctx->dm_coefficients) DMDestroy(&level_ctx->dm_coefficients);
436:   if (level_ctx->dm_faces) DMDestroy(&level_ctx->dm_faces);
437:   if (level_ctx->coeff) VecDestroy(&level_ctx->coeff);
438:   PetscFree(*p_level_ctx);
439:   return 0;
440: }

442: static PetscErrorCode CtxCreateAndSetFromOptions(Ctx *p_ctx)
443: {
444:   Ctx ctx;

447:   PetscMalloc1(1, p_ctx);
448:   ctx = *p_ctx;

450:   ctx->comm         = PETSC_COMM_WORLD;
451:   ctx->pin_pressure = PETSC_FALSE;
452:   PetscOptionsGetBool(NULL, NULL, "-pin_pressure", &ctx->pin_pressure, NULL);
453:   ctx->dim = 3;
454:   PetscOptionsGetInt(NULL, NULL, "-dim", &ctx->dim, NULL);
455:   if (ctx->dim <= 2) {
456:     ctx->cells_x = 32;
457:   } else {
458:     ctx->cells_x = 16;
459:   }
460:   PetscOptionsGetInt(NULL, NULL, "-s", &ctx->cells_x, NULL); /* shortcut. Usually, use -stag_grid_x etc. */
461:   ctx->cells_z = ctx->cells_y = ctx->cells_x;
462:   ctx->xmin = ctx->ymin = ctx->zmin = 0.0;
463:   {
464:     PetscBool nondimensional = PETSC_TRUE;

466:     PetscOptionsGetBool(NULL, NULL, "-nondimensional", &nondimensional, NULL);
467:     if (nondimensional) {
468:       ctx->xmax = ctx->ymax = ctx->zmax = 1.0;
469:       ctx->rho1                         = 0.0;
470:       ctx->rho2                         = 1.0;
471:       ctx->eta1                         = 1.0;
472:       ctx->eta2                         = 1e2;
473:       ctx->gy                           = -1.0; /* downwards */
474:     } else {
475:       ctx->xmax = 1e6;
476:       ctx->ymax = 1.5e6;
477:       ctx->zmax = 1e6;
478:       ctx->rho1 = 3200;
479:       ctx->rho2 = 3300;
480:       ctx->eta1 = 1e20;
481:       ctx->eta2 = 1e22;
482:       ctx->gy   = -10.0; /* downwards */
483:     }
484:   }
485:   {
486:     PetscBool isoviscous;

488:     isoviscous = PETSC_FALSE;
489:     PetscOptionsGetScalar(NULL, NULL, "-eta1", &ctx->eta1, NULL);
490:     PetscOptionsGetBool(NULL, NULL, "-isoviscous", &isoviscous, NULL);
491:     if (isoviscous) {
492:       ctx->eta2   = ctx->eta1;
493:       ctx->GetEta = GetEta_constant; /* override */
494:     } else {
495:       PetscOptionsGetScalar(NULL, NULL, "-eta2", &ctx->eta2, NULL);
496:     }
497:   }
498:   {
499:     char      mode[1024] = "sinker";
500:     PetscBool is_layers, is_blob, is_sinker_box, is_sinker_sphere;

502:     PetscOptionsGetString(NULL, NULL, "-coefficients", mode, sizeof(mode), NULL);
503:     PetscStrncmp(mode, "layers", sizeof(mode), &is_layers);
504:     PetscStrncmp(mode, "sinker", sizeof(mode), &is_sinker_box);
505:     if (!is_sinker_box) PetscStrncmp(mode, "sinker_box", sizeof(mode), &is_sinker_box);
506:     PetscStrncmp(mode, "sinker_sphere", sizeof(mode), &is_sinker_sphere);
507:     PetscStrncmp(mode, "blob", sizeof(mode), &is_blob);

509:     if (is_layers) {
510:       ctx->GetRho = GetRho_layers;
511:       ctx->GetEta = GetEta_layers;
512:     }
513:     if (is_blob) {
514:       if (ctx->dim == 3) {
515:         ctx->GetRho = GetRho_blob3;
516:         ctx->GetEta = GetEta_blob3;
517:       } else SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_SUP, "Not implemented for dimension %" PetscInt_FMT, ctx->dim);
518:     }
519:     if (is_sinker_box) {
520:       if (ctx->dim == 2) {
521:         ctx->GetRho = GetRho_sinker_box2;
522:         ctx->GetEta = GetEta_sinker_box2;
523:       } else if (ctx->dim == 3) {
524:         ctx->GetRho = GetRho_sinker_box3;
525:         ctx->GetEta = GetEta_sinker_box3;
526:       } else SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_SUP, "Not implemented for dimension %" PetscInt_FMT, ctx->dim);
527:     }
528:     if (is_sinker_sphere) {
529:       if (ctx->dim == 3) {
530:         ctx->GetRho = GetRho_sinker_sphere3;
531:         ctx->GetEta = GetEta_sinker_sphere3;
532:       } else SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_SUP, "Not implemented for dimension %" PetscInt_FMT, ctx->dim);
533:     }
534:   }

536:   /* Per-level data */
537:   ctx->n_levels = 1;
538:   PetscOptionsGetInt(NULL, NULL, "-levels", &ctx->n_levels, NULL);
539:   PetscMalloc1(ctx->n_levels, &ctx->levels);
540:   for (PetscInt i = 0; i < ctx->n_levels; ++i) LevelCtxCreate(&ctx->levels[i]);
541:   return 0;
542: }

544: static PetscErrorCode CtxDestroy(Ctx *p_ctx)
545: {
546:   Ctx ctx;

549:   ctx = *p_ctx;
550:   for (PetscInt i = 0; i < ctx->n_levels; ++i) LevelCtxDestroy(&ctx->levels[i]);
551:   PetscFree(ctx->levels);
552:   PetscFree(*p_ctx);
553:   return 0;
554: }

556: static PetscErrorCode SystemParametersCreate(SystemParameters *parameters, Ctx ctx)
557: {
559:   PetscMalloc1(1, parameters);
560:   (*parameters)->ctx                  = ctx;
561:   (*parameters)->level                = ctx->n_levels - 1;
562:   (*parameters)->include_inverse_visc = PETSC_FALSE;
563:   (*parameters)->faces_only           = PETSC_FALSE;
564:   return 0;
565: }

567: static PetscErrorCode SystemParametersDestroy(SystemParameters *parameters)
568: {
570:   PetscFree(*parameters);
571:   *parameters = NULL;
572:   return 0;
573: }

575: static PetscErrorCode CreateSystem2d(SystemParameters parameters, Mat *pA, Vec *pRhs)
576: {
577:   PetscInt    N[2];
578:   PetscInt    ex, ey, startx, starty, nx, ny;
579:   Mat         A;
580:   Vec         rhs;
581:   PetscReal   hx, hy, dv;
582:   Vec         coefficients_local;
583:   PetscBool   build_rhs;
584:   DM          dm_main, dm_coefficients;
585:   PetscScalar K_cont, K_bound;
586:   Ctx         ctx   = parameters->ctx;
587:   PetscInt    level = parameters->level;

590:   if (parameters->faces_only) {
591:     dm_main = ctx->levels[level]->dm_faces;
592:   } else {
593:     dm_main = ctx->levels[level]->dm_stokes;
594:   }
595:   dm_coefficients = ctx->levels[level]->dm_coefficients;
596:   K_cont          = ctx->levels[level]->K_cont;
597:   K_bound         = ctx->levels[level]->K_bound;
598:   DMCreateMatrix(dm_main, pA);
599:   A         = *pA;
600:   build_rhs = (PetscBool)(pRhs != NULL);
602:   if (build_rhs) {
603:     DMCreateGlobalVector(dm_main, pRhs);
604:     rhs = *pRhs;
605:   } else {
606:     rhs = NULL;
607:   }
608:   DMStagGetCorners(dm_main, &startx, &starty, NULL, &nx, &ny, NULL, NULL, NULL, NULL);
609:   DMStagGetGlobalSizes(dm_main, &N[0], &N[1], NULL);
610:   hx = ctx->levels[level]->hx_characteristic;
611:   hy = ctx->levels[level]->hy_characteristic;
612:   dv = hx * hy;
613:   DMGetLocalVector(dm_coefficients, &coefficients_local);
614:   DMGlobalToLocal(dm_coefficients, ctx->levels[level]->coeff, INSERT_VALUES, coefficients_local);

616:   /* Loop over all local elements */
617:   for (ey = starty; ey < starty + ny; ++ey) { /* With DMStag, always iterate x fastest, y second fastest, z slowest */
618:     for (ex = startx; ex < startx + nx; ++ex) {
619:       const PetscBool left_boundary   = (PetscBool)(ex == 0);
620:       const PetscBool right_boundary  = (PetscBool)(ex == N[0] - 1);
621:       const PetscBool bottom_boundary = (PetscBool)(ey == 0);
622:       const PetscBool top_boundary    = (PetscBool)(ey == N[1] - 1);

624:       if (ey == N[1] - 1) {
625:         /* Top boundary velocity Dirichlet */
626:         DMStagStencil     row;
627:         const PetscScalar val_rhs = 0.0;
628:         const PetscScalar val_A   = K_bound;

630:         row.i   = ex;
631:         row.j   = ey;
632:         row.loc = DMSTAG_UP;
633:         row.c   = 0;
634:         DMStagMatSetValuesStencil(dm_main, A, 1, &row, 1, &row, &val_A, INSERT_VALUES);
635:         if (build_rhs) DMStagVecSetValuesStencil(dm_main, rhs, 1, &row, &val_rhs, INSERT_VALUES);
636:       }

638:       if (ey == 0) {
639:         /* Bottom boundary velocity Dirichlet */
640:         DMStagStencil     row;
641:         const PetscScalar val_rhs = 0.0;
642:         const PetscScalar val_A   = K_bound;

644:         row.i   = ex;
645:         row.j   = ey;
646:         row.loc = DMSTAG_DOWN;
647:         row.c   = 0;
648:         DMStagMatSetValuesStencil(dm_main, A, 1, &row, 1, &row, &val_A, INSERT_VALUES);
649:         if (build_rhs) DMStagVecSetValuesStencil(dm_main, rhs, 1, &row, &val_rhs, INSERT_VALUES);
650:       } else {
651:         /* Y-momentum equation : (u_xx + u_yy) - p_y = f^y
652:            includes non-zero forcing and free-slip boundary conditions */
653:         PetscInt      count;
654:         DMStagStencil row, col[11];
655:         PetscScalar   val_A[11];
656:         DMStagStencil rhoPoint[2];
657:         PetscScalar   rho[2], val_rhs;
658:         DMStagStencil etaPoint[4];
659:         PetscScalar   eta[4], eta_left, eta_right, eta_up, eta_down;

661:         row.i   = ex;
662:         row.j   = ey;
663:         row.loc = DMSTAG_DOWN;
664:         row.c   = 0;

666:         /* get rho values  and compute rhs value*/
667:         rhoPoint[0].i   = ex;
668:         rhoPoint[0].j   = ey;
669:         rhoPoint[0].loc = DMSTAG_DOWN_LEFT;
670:         rhoPoint[0].c   = 1;
671:         rhoPoint[1].i   = ex;
672:         rhoPoint[1].j   = ey;
673:         rhoPoint[1].loc = DMSTAG_DOWN_RIGHT;
674:         rhoPoint[1].c   = 1;
675:         DMStagVecGetValuesStencil(dm_coefficients, coefficients_local, 2, rhoPoint, rho);
676:         val_rhs = -ctx->gy * dv * 0.5 * (rho[0] + rho[1]);

678:         /* Get eta values */
679:         etaPoint[0].i   = ex;
680:         etaPoint[0].j   = ey;
681:         etaPoint[0].loc = DMSTAG_DOWN_LEFT;
682:         etaPoint[0].c   = 0; /* Left  */
683:         etaPoint[1].i   = ex;
684:         etaPoint[1].j   = ey;
685:         etaPoint[1].loc = DMSTAG_DOWN_RIGHT;
686:         etaPoint[1].c   = 0; /* Right */
687:         etaPoint[2].i   = ex;
688:         etaPoint[2].j   = ey;
689:         etaPoint[2].loc = DMSTAG_ELEMENT;
690:         etaPoint[2].c   = 0; /* Up    */
691:         etaPoint[3].i   = ex;
692:         etaPoint[3].j   = ey - 1;
693:         etaPoint[3].loc = DMSTAG_ELEMENT;
694:         etaPoint[3].c   = 0; /* Down  */
695:         DMStagVecGetValuesStencil(dm_coefficients, coefficients_local, 4, etaPoint, eta);
696:         eta_left  = eta[0];
697:         eta_right = eta[1];
698:         eta_up    = eta[2];
699:         eta_down  = eta[3];

701:         count = 0;

703:         col[count]   = row;
704:         val_A[count] = -2.0 * dv * (eta_down + eta_up) / (hy * hy);
705:         if (!left_boundary) val_A[count] += -1.0 * dv * eta_left / (hx * hx);
706:         if (!right_boundary) val_A[count] += -1.0 * dv * eta_right / (hx * hx);
707:         ++count;

709:         col[count].i   = ex;
710:         col[count].j   = ey - 1;
711:         col[count].loc = DMSTAG_DOWN;
712:         col[count].c   = 0;
713:         val_A[count]   = 2.0 * dv * eta_down / (hy * hy);
714:         ++count;
715:         col[count].i   = ex;
716:         col[count].j   = ey + 1;
717:         col[count].loc = DMSTAG_DOWN;
718:         col[count].c   = 0;
719:         val_A[count]   = 2.0 * dv * eta_up / (hy * hy);
720:         ++count;
721:         if (!left_boundary) {
722:           col[count].i   = ex - 1;
723:           col[count].j   = ey;
724:           col[count].loc = DMSTAG_DOWN;
725:           col[count].c   = 0;
726:           val_A[count]   = dv * eta_left / (hx * hx);
727:           ++count;
728:         }
729:         if (!right_boundary) {
730:           col[count].i   = ex + 1;
731:           col[count].j   = ey;
732:           col[count].loc = DMSTAG_DOWN;
733:           col[count].c   = 0;
734:           val_A[count]   = dv * eta_right / (hx * hx);
735:           ++count;
736:         }
737:         col[count].i   = ex;
738:         col[count].j   = ey - 1;
739:         col[count].loc = DMSTAG_LEFT;
740:         col[count].c   = 0;
741:         val_A[count]   = dv * eta_left / (hx * hy);
742:         ++count; /* down left x edge */
743:         col[count].i   = ex;
744:         col[count].j   = ey - 1;
745:         col[count].loc = DMSTAG_RIGHT;
746:         col[count].c   = 0;
747:         val_A[count]   = -1.0 * dv * eta_right / (hx * hy);
748:         ++count; /* down right x edge */
749:         col[count].i   = ex;
750:         col[count].j   = ey;
751:         col[count].loc = DMSTAG_LEFT;
752:         col[count].c   = 0;
753:         val_A[count]   = -1.0 * dv * eta_left / (hx * hy);
754:         ++count; /* up left x edge */
755:         col[count].i   = ex;
756:         col[count].j   = ey;
757:         col[count].loc = DMSTAG_RIGHT;
758:         col[count].c   = 0;
759:         val_A[count]   = dv * eta_right / (hx * hy);
760:         ++count; /* up right x edge */
761:         if (!parameters->faces_only) {
762:           col[count].i   = ex;
763:           col[count].j   = ey - 1;
764:           col[count].loc = DMSTAG_ELEMENT;
765:           col[count].c   = 0;
766:           val_A[count]   = K_cont * dv / hy;
767:           ++count;
768:           col[count].i   = ex;
769:           col[count].j   = ey;
770:           col[count].loc = DMSTAG_ELEMENT;
771:           col[count].c   = 0;
772:           val_A[count]   = -1.0 * K_cont * dv / hy;
773:           ++count;
774:         }

776:         DMStagMatSetValuesStencil(dm_main, A, 1, &row, count, col, val_A, INSERT_VALUES);
777:         if (build_rhs) DMStagVecSetValuesStencil(dm_main, rhs, 1, &row, &val_rhs, INSERT_VALUES);
778:       }

780:       if (ex == N[0] - 1) {
781:         /* Right Boundary velocity Dirichlet */
782:         /* Redundant in the corner */
783:         DMStagStencil     row;
784:         const PetscScalar val_rhs = 0.0;
785:         const PetscScalar val_A   = K_bound;

787:         row.i   = ex;
788:         row.j   = ey;
789:         row.loc = DMSTAG_RIGHT;
790:         row.c   = 0;
791:         DMStagMatSetValuesStencil(dm_main, A, 1, &row, 1, &row, &val_A, INSERT_VALUES);
792:         if (build_rhs) DMStagVecSetValuesStencil(dm_main, rhs, 1, &row, &val_rhs, INSERT_VALUES);
793:       }
794:       if (ex == 0) {
795:         /* Left velocity Dirichlet */
796:         DMStagStencil     row;
797:         const PetscScalar val_rhs = 0.0;
798:         const PetscScalar val_A   = K_bound;

800:         row.i   = ex;
801:         row.j   = ey;
802:         row.loc = DMSTAG_LEFT;
803:         row.c   = 0;
804:         DMStagMatSetValuesStencil(dm_main, A, 1, &row, 1, &row, &val_A, INSERT_VALUES);
805:         if (build_rhs) DMStagVecSetValuesStencil(dm_main, rhs, 1, &row, &val_rhs, INSERT_VALUES);
806:       } else {
807:         /* X-momentum equation : (u_xx + u_yy) - p_x = f^x
808:           zero RHS, including free-slip boundary conditions */
809:         PetscInt          count;
810:         DMStagStencil     row, col[11];
811:         PetscScalar       val_A[11];
812:         DMStagStencil     etaPoint[4];
813:         PetscScalar       eta[4], eta_left, eta_right, eta_up, eta_down;
814:         const PetscScalar val_rhs = 0.0;

816:         row.i   = ex;
817:         row.j   = ey;
818:         row.loc = DMSTAG_LEFT;
819:         row.c   = 0;

821:         /* Get eta values */
822:         etaPoint[0].i   = ex - 1;
823:         etaPoint[0].j   = ey;
824:         etaPoint[0].loc = DMSTAG_ELEMENT;
825:         etaPoint[0].c   = 0; /* Left  */
826:         etaPoint[1].i   = ex;
827:         etaPoint[1].j   = ey;
828:         etaPoint[1].loc = DMSTAG_ELEMENT;
829:         etaPoint[1].c   = 0; /* Right */
830:         etaPoint[2].i   = ex;
831:         etaPoint[2].j   = ey;
832:         etaPoint[2].loc = DMSTAG_UP_LEFT;
833:         etaPoint[2].c   = 0; /* Up    */
834:         etaPoint[3].i   = ex;
835:         etaPoint[3].j   = ey;
836:         etaPoint[3].loc = DMSTAG_DOWN_LEFT;
837:         etaPoint[3].c   = 0; /* Down  */
838:         DMStagVecGetValuesStencil(dm_coefficients, coefficients_local, 4, etaPoint, eta);
839:         eta_left  = eta[0];
840:         eta_right = eta[1];
841:         eta_up    = eta[2];
842:         eta_down  = eta[3];

844:         count        = 0;
845:         col[count]   = row;
846:         val_A[count] = -2.0 * dv * (eta_left + eta_right) / (hx * hx);
847:         if (!bottom_boundary) val_A[count] += -1.0 * dv * eta_down / (hy * hy);
848:         if (!top_boundary) val_A[count] += -1.0 * dv * eta_up / (hy * hy);
849:         ++count;

851:         if (!bottom_boundary) {
852:           col[count].i   = ex;
853:           col[count].j   = ey - 1;
854:           col[count].loc = DMSTAG_LEFT;
855:           col[count].c   = 0;
856:           val_A[count]   = dv * eta_down / (hy * hy);
857:           ++count;
858:         }
859:         if (!top_boundary) {
860:           col[count].i   = ex;
861:           col[count].j   = ey + 1;
862:           col[count].loc = DMSTAG_LEFT;
863:           col[count].c   = 0;
864:           val_A[count]   = dv * eta_up / (hy * hy);
865:           ++count;
866:         }
867:         col[count].i   = ex - 1;
868:         col[count].j   = ey;
869:         col[count].loc = DMSTAG_LEFT;
870:         col[count].c   = 0;
871:         val_A[count]   = 2.0 * dv * eta_left / (hx * hx);
872:         ++count;
873:         col[count].i   = ex + 1;
874:         col[count].j   = ey;
875:         col[count].loc = DMSTAG_LEFT;
876:         col[count].c   = 0;
877:         val_A[count]   = 2.0 * dv * eta_right / (hx * hx);
878:         ++count;
879:         col[count].i   = ex - 1;
880:         col[count].j   = ey;
881:         col[count].loc = DMSTAG_DOWN;
882:         col[count].c   = 0;
883:         val_A[count]   = dv * eta_down / (hx * hy);
884:         ++count; /* down left */
885:         col[count].i   = ex;
886:         col[count].j   = ey;
887:         col[count].loc = DMSTAG_DOWN;
888:         col[count].c   = 0;
889:         val_A[count]   = -1.0 * dv * eta_down / (hx * hy);
890:         ++count; /* down right */
891:         col[count].i   = ex - 1;
892:         col[count].j   = ey;
893:         col[count].loc = DMSTAG_UP;
894:         col[count].c   = 0;
895:         val_A[count]   = -1.0 * dv * eta_up / (hx * hy);
896:         ++count; /* up left */
897:         col[count].i   = ex;
898:         col[count].j   = ey;
899:         col[count].loc = DMSTAG_UP;
900:         col[count].c   = 0;
901:         val_A[count]   = dv * eta_up / (hx * hy);
902:         ++count; /* up right */
903:         if (!parameters->faces_only) {
904:           col[count].i   = ex - 1;
905:           col[count].j   = ey;
906:           col[count].loc = DMSTAG_ELEMENT;
907:           col[count].c   = 0;
908:           val_A[count]   = K_cont * dv / hx;
909:           ++count;
910:           col[count].i   = ex;
911:           col[count].j   = ey;
912:           col[count].loc = DMSTAG_ELEMENT;
913:           col[count].c   = 0;
914:           val_A[count]   = -1.0 * K_cont * dv / hx;
915:           ++count;
916:         }

918:         DMStagMatSetValuesStencil(dm_main, A, 1, &row, count, col, val_A, INSERT_VALUES);
919:         if (build_rhs) DMStagVecSetValuesStencil(dm_main, rhs, 1, &row, &val_rhs, INSERT_VALUES);
920:       }

922:       /* P equation : u_x + v_y = 0

924:          Note that this includes an explicit zero on the diagonal. This is only needed for
925:          direct solvers (not required if using an iterative solver and setting the constant-pressure nullspace)

927:         Note: the scaling by dv is not chosen in a principled way and is likely sub-optimal
928:        */
929:       if (!parameters->faces_only) {
930:         if (ctx->pin_pressure && ex == 0 && ey == 0) { /* Pin the first pressure node to zero, if requested */
931:           DMStagStencil     row;
932:           const PetscScalar val_A   = K_bound;
933:           const PetscScalar val_rhs = 0.0;

935:           row.i   = ex;
936:           row.j   = ey;
937:           row.loc = DMSTAG_ELEMENT;
938:           row.c   = 0;
939:           DMStagMatSetValuesStencil(dm_main, A, 1, &row, 1, &row, &val_A, INSERT_VALUES);
940:           if (build_rhs) DMStagVecSetValuesStencil(dm_main, rhs, 1, &row, &val_rhs, INSERT_VALUES);
941:         } else {
942:           DMStagStencil     row, col[5];
943:           PetscScalar       val_A[5];
944:           const PetscScalar val_rhs = 0.0;

946:           row.i      = ex;
947:           row.j      = ey;
948:           row.loc    = DMSTAG_ELEMENT;
949:           row.c      = 0;
950:           col[0].i   = ex;
951:           col[0].j   = ey;
952:           col[0].loc = DMSTAG_LEFT;
953:           col[0].c   = 0;
954:           val_A[0]   = -1.0 * K_cont * dv / hx;
955:           col[1].i   = ex;
956:           col[1].j   = ey;
957:           col[1].loc = DMSTAG_RIGHT;
958:           col[1].c   = 0;
959:           val_A[1]   = K_cont * dv / hx;
960:           col[2].i   = ex;
961:           col[2].j   = ey;
962:           col[2].loc = DMSTAG_DOWN;
963:           col[2].c   = 0;
964:           val_A[2]   = -1.0 * K_cont * dv / hy;
965:           col[3].i   = ex;
966:           col[3].j   = ey;
967:           col[3].loc = DMSTAG_UP;
968:           col[3].c   = 0;
969:           val_A[3]   = K_cont * dv / hy;
970:           col[4]     = row;
971:           val_A[4]   = 0.0;
972:           DMStagMatSetValuesStencil(dm_main, A, 1, &row, 5, col, val_A, INSERT_VALUES);
973:           if (build_rhs) DMStagVecSetValuesStencil(dm_main, rhs, 1, &row, &val_rhs, INSERT_VALUES);
974:         }
975:       }
976:     }
977:   }
978:   DMRestoreLocalVector(dm_coefficients, &coefficients_local);

980:   /* Add additional inverse viscosity terms (for use in building a preconditioning matrix) */
981:   if (parameters->include_inverse_visc) {
983:     OperatorInsertInverseViscosityPressureTerms(dm_main, dm_coefficients, ctx->levels[level]->coeff, 1.0, A);
984:   }

986:   MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY);
987:   if (build_rhs) VecAssemblyBegin(rhs);
988:   MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY);
989:   if (build_rhs) VecAssemblyEnd(rhs);
990:   return 0;
991: }

993: static PetscErrorCode CreateSystem3d(SystemParameters parameters, Mat *pA, Vec *pRhs)
994: {
995:   PetscInt    N[3];
996:   PetscInt    ex, ey, ez, startx, starty, startz, nx, ny, nz;
997:   Mat         A;
998:   PetscReal   hx, hy, hz, dv;
999:   PetscInt    pinx, piny, pinz;
1000:   Vec         coeff_local, rhs;
1001:   PetscBool   build_rhs;
1002:   DM          dm_main, dm_coefficients;
1003:   PetscScalar K_cont, K_bound;
1004:   Ctx         ctx   = parameters->ctx;
1005:   PetscInt    level = parameters->level;

1008:   if (parameters->faces_only) {
1009:     dm_main = ctx->levels[level]->dm_faces;
1010:   } else {
1011:     dm_main = ctx->levels[level]->dm_stokes;
1012:   }
1013:   dm_coefficients = ctx->levels[level]->dm_coefficients;
1014:   K_cont          = ctx->levels[level]->K_cont;
1015:   K_bound         = ctx->levels[level]->K_bound;
1016:   DMCreateMatrix(dm_main, pA);
1017:   A         = *pA;
1018:   build_rhs = (PetscBool)(pRhs != NULL);
1019:   if (build_rhs) {
1020:     DMCreateGlobalVector(dm_main, pRhs);
1021:     rhs = *pRhs;
1022:   } else {
1023:     rhs = NULL;
1024:   }
1025:   DMStagGetCorners(dm_main, &startx, &starty, &startz, &nx, &ny, &nz, NULL, NULL, NULL);
1026:   DMStagGetGlobalSizes(dm_main, &N[0], &N[1], &N[2]);
1027:   hx = ctx->levels[level]->hx_characteristic;
1028:   hy = ctx->levels[level]->hy_characteristic;
1029:   hz = ctx->levels[level]->hz_characteristic;
1030:   dv = hx * hy * hz;
1031:   DMGetLocalVector(dm_coefficients, &coeff_local);
1032:   DMGlobalToLocal(dm_coefficients, ctx->levels[level]->coeff, INSERT_VALUES, coeff_local);

1035:   pinx = 1;
1036:   piny = 0;
1037:   pinz = 0; /* Depends on assertion above that there are at least two element in the x direction */

1039:   /* Loop over all local elements.

1041:      For each element, fill 4-7 rows of the matrix, corresponding to
1042:      - the pressure degree of freedom (dof), centered on the element
1043:      - the 3 velocity dofs on left, bottom, and back faces of the element
1044:      - velocity dof on the right, top, and front faces of the element (only on domain boundaries)

1046:    */
1047:   for (ez = startz; ez < startz + nz; ++ez) {
1048:     for (ey = starty; ey < starty + ny; ++ey) {
1049:       for (ex = startx; ex < startx + nx; ++ex) {
1050:         const PetscBool left_boundary   = (PetscBool)(ex == 0);
1051:         const PetscBool right_boundary  = (PetscBool)(ex == N[0] - 1);
1052:         const PetscBool bottom_boundary = (PetscBool)(ey == 0);
1053:         const PetscBool top_boundary    = (PetscBool)(ey == N[1] - 1);
1054:         const PetscBool back_boundary   = (PetscBool)(ez == 0);
1055:         const PetscBool front_boundary  = (PetscBool)(ez == N[2] - 1);

1057:         /* Note that below, we depend on the check above that there is never one
1058:            element (globally) in a given direction.  Thus, for example, an
1059:            element is never both on the left and right boundary */

1061:         /* X-faces - right boundary */
1062:         if (right_boundary) {
1063:           /* Right x-velocity Dirichlet */
1064:           DMStagStencil     row;
1065:           const PetscScalar val_rhs = 0.0;
1066:           const PetscScalar val_A   = K_bound;

1068:           row.i   = ex;
1069:           row.j   = ey;
1070:           row.k   = ez;
1071:           row.loc = DMSTAG_RIGHT;
1072:           row.c   = 0;
1073:           DMStagMatSetValuesStencil(dm_main, A, 1, &row, 1, &row, &val_A, INSERT_VALUES);
1074:           if (build_rhs) DMStagVecSetValuesStencil(dm_main, rhs, 1, &row, &val_rhs, INSERT_VALUES);
1075:         }

1077:         /* X faces - left*/
1078:         {
1079:           DMStagStencil row;

1081:           row.i   = ex;
1082:           row.j   = ey;
1083:           row.k   = ez;
1084:           row.loc = DMSTAG_LEFT;
1085:           row.c   = 0;

1087:           if (left_boundary) {
1088:             /* Left x-velocity Dirichlet */
1089:             const PetscScalar val_rhs = 0.0;
1090:             const PetscScalar val_A   = K_bound;

1092:             DMStagMatSetValuesStencil(dm_main, A, 1, &row, 1, &row, &val_A, INSERT_VALUES);
1093:             if (build_rhs) DMStagVecSetValuesStencil(dm_main, rhs, 1, &row, &val_rhs, INSERT_VALUES);
1094:           } else {
1095:             /* X-momentum equation */
1096:             PetscInt          count;
1097:             DMStagStencil     col[17];
1098:             PetscScalar       val_A[17];
1099:             DMStagStencil     eta_point[6];
1100:             PetscScalar       eta[6], eta_left, eta_right, eta_up, eta_down, eta_back, eta_front; /* relative to the left face */
1101:             const PetscScalar val_rhs = 0.0;

1103:             /* Get eta values */
1104:             eta_point[0].i   = ex - 1;
1105:             eta_point[0].j   = ey;
1106:             eta_point[0].k   = ez;
1107:             eta_point[0].loc = DMSTAG_ELEMENT;
1108:             eta_point[0].c   = 0; /* Left  */
1109:             eta_point[1].i   = ex;
1110:             eta_point[1].j   = ey;
1111:             eta_point[1].k   = ez;
1112:             eta_point[1].loc = DMSTAG_ELEMENT;
1113:             eta_point[1].c   = 0; /* Right */
1114:             eta_point[2].i   = ex;
1115:             eta_point[2].j   = ey;
1116:             eta_point[2].k   = ez;
1117:             eta_point[2].loc = DMSTAG_DOWN_LEFT;
1118:             eta_point[2].c   = 0; /* Down  */
1119:             eta_point[3].i   = ex;
1120:             eta_point[3].j   = ey;
1121:             eta_point[3].k   = ez;
1122:             eta_point[3].loc = DMSTAG_UP_LEFT;
1123:             eta_point[3].c   = 0; /* Up    */
1124:             eta_point[4].i   = ex;
1125:             eta_point[4].j   = ey;
1126:             eta_point[4].k   = ez;
1127:             eta_point[4].loc = DMSTAG_BACK_LEFT;
1128:             eta_point[4].c   = 0; /* Back  */
1129:             eta_point[5].i   = ex;
1130:             eta_point[5].j   = ey;
1131:             eta_point[5].k   = ez;
1132:             eta_point[5].loc = DMSTAG_FRONT_LEFT;
1133:             eta_point[5].c   = 0; /* Front  */
1134:             DMStagVecGetValuesStencil(dm_coefficients, coeff_local, 6, eta_point, eta);
1135:             eta_left  = eta[0];
1136:             eta_right = eta[1];
1137:             eta_down  = eta[2];
1138:             eta_up    = eta[3];
1139:             eta_back  = eta[4];
1140:             eta_front = eta[5];

1142:             count = 0;

1144:             col[count]   = row;
1145:             val_A[count] = -2.0 * dv * (eta_left + eta_right) / (hx * hx);
1146:             if (!top_boundary) val_A[count] += -1.0 * dv * eta_up / (hy * hy);
1147:             if (!bottom_boundary) val_A[count] += -1.0 * dv * eta_down / (hy * hy);
1148:             if (!back_boundary) val_A[count] += -1.0 * dv * eta_back / (hz * hz);
1149:             if (!front_boundary) val_A[count] += -1.0 * dv * eta_front / (hz * hz);
1150:             ++count;

1152:             col[count].i   = ex - 1;
1153:             col[count].j   = ey;
1154:             col[count].k   = ez;
1155:             col[count].loc = DMSTAG_LEFT;
1156:             col[count].c   = 0;
1157:             val_A[count]   = 2.0 * dv * eta_left / (hx * hx);
1158:             ++count;
1159:             col[count].i   = ex + 1;
1160:             col[count].j   = ey;
1161:             col[count].k   = ez;
1162:             col[count].loc = DMSTAG_LEFT;
1163:             col[count].c   = 0;
1164:             val_A[count]   = 2.0 * dv * eta_right / (hx * hx);
1165:             ++count;
1166:             if (!bottom_boundary) {
1167:               col[count].i   = ex;
1168:               col[count].j   = ey - 1;
1169:               col[count].k   = ez;
1170:               col[count].loc = DMSTAG_LEFT;
1171:               col[count].c   = 0;
1172:               val_A[count]   = dv * eta_down / (hy * hy);
1173:               ++count;
1174:             }
1175:             if (!top_boundary) {
1176:               col[count].i   = ex;
1177:               col[count].j   = ey + 1;
1178:               col[count].k   = ez;
1179:               col[count].loc = DMSTAG_LEFT;
1180:               col[count].c   = 0;
1181:               val_A[count]   = dv * eta_up / (hy * hy);
1182:               ++count;
1183:             }
1184:             if (!back_boundary) {
1185:               col[count].i   = ex;
1186:               col[count].j   = ey;
1187:               col[count].k   = ez - 1;
1188:               col[count].loc = DMSTAG_LEFT;
1189:               col[count].c   = 0;
1190:               val_A[count]   = dv * eta_back / (hz * hz);
1191:               ++count;
1192:             }
1193:             if (!front_boundary) {
1194:               col[count].i   = ex;
1195:               col[count].j   = ey;
1196:               col[count].k   = ez + 1;
1197:               col[count].loc = DMSTAG_LEFT;
1198:               col[count].c   = 0;
1199:               val_A[count]   = dv * eta_front / (hz * hz);
1200:               ++count;
1201:             }

1203:             col[count].i   = ex - 1;
1204:             col[count].j   = ey;
1205:             col[count].k   = ez;
1206:             col[count].loc = DMSTAG_DOWN;
1207:             col[count].c   = 0;
1208:             val_A[count]   = dv * eta_down / (hx * hy);
1209:             ++count; /* down left */
1210:             col[count].i   = ex;
1211:             col[count].j   = ey;
1212:             col[count].k   = ez;
1213:             col[count].loc = DMSTAG_DOWN;
1214:             col[count].c   = 0;
1215:             val_A[count]   = -1.0 * dv * eta_down / (hx * hy);
1216:             ++count; /* down right */

1218:             col[count].i   = ex - 1;
1219:             col[count].j   = ey;
1220:             col[count].k   = ez;
1221:             col[count].loc = DMSTAG_UP;
1222:             col[count].c   = 0;
1223:             val_A[count]   = -1.0 * dv * eta_up / (hx * hy);
1224:             ++count; /* up left */
1225:             col[count].i   = ex;
1226:             col[count].j   = ey;
1227:             col[count].k   = ez;
1228:             col[count].loc = DMSTAG_UP;
1229:             col[count].c   = 0;
1230:             val_A[count]   = dv * eta_up / (hx * hy);
1231:             ++count; /* up right */

1233:             col[count].i   = ex - 1;
1234:             col[count].j   = ey;
1235:             col[count].k   = ez;
1236:             col[count].loc = DMSTAG_BACK;
1237:             col[count].c   = 0;
1238:             val_A[count]   = dv * eta_back / (hx * hz);
1239:             ++count; /* back left */
1240:             col[count].i   = ex;
1241:             col[count].j   = ey;
1242:             col[count].k   = ez;
1243:             col[count].loc = DMSTAG_BACK;
1244:             col[count].c   = 0;
1245:             val_A[count]   = -1.0 * dv * eta_back / (hx * hz);
1246:             ++count; /* back right */

1248:             col[count].i   = ex - 1;
1249:             col[count].j   = ey;
1250:             col[count].k   = ez;
1251:             col[count].loc = DMSTAG_FRONT;
1252:             col[count].c   = 0;
1253:             val_A[count]   = -1.0 * dv * eta_front / (hx * hz);
1254:             ++count; /* front left */
1255:             col[count].i   = ex;
1256:             col[count].j   = ey;
1257:             col[count].k   = ez;
1258:             col[count].loc = DMSTAG_FRONT;
1259:             col[count].c   = 0;
1260:             val_A[count]   = dv * eta_front / (hx * hz);
1261:             ++count; /* front right */

1263:             if (!parameters->faces_only) {
1264:               col[count].i   = ex - 1;
1265:               col[count].j   = ey;
1266:               col[count].k   = ez;
1267:               col[count].loc = DMSTAG_ELEMENT;
1268:               col[count].c   = 0;
1269:               val_A[count]   = K_cont * dv / hx;
1270:               ++count;
1271:               col[count].i   = ex;
1272:               col[count].j   = ey;
1273:               col[count].k   = ez;
1274:               col[count].loc = DMSTAG_ELEMENT;
1275:               col[count].c   = 0;
1276:               val_A[count]   = -1.0 * K_cont * dv / hx;
1277:               ++count;
1278:             }

1280:             DMStagMatSetValuesStencil(dm_main, A, 1, &row, count, col, val_A, INSERT_VALUES);
1281:             if (build_rhs) DMStagVecSetValuesStencil(dm_main, rhs, 1, &row, &val_rhs, INSERT_VALUES);
1282:           }
1283:         }

1285:         /* Y faces - top boundary */
1286:         if (top_boundary) {
1287:           /* Top y-velocity Dirichlet */
1288:           DMStagStencil     row;
1289:           const PetscScalar val_rhs = 0.0;
1290:           const PetscScalar val_A   = K_bound;

1292:           row.i   = ex;
1293:           row.j   = ey;
1294:           row.k   = ez;
1295:           row.loc = DMSTAG_UP;
1296:           row.c   = 0;
1297:           DMStagMatSetValuesStencil(dm_main, A, 1, &row, 1, &row, &val_A, INSERT_VALUES);
1298:           if (build_rhs) DMStagVecSetValuesStencil(dm_main, rhs, 1, &row, &val_rhs, INSERT_VALUES);
1299:         }

1301:         /* Y faces - down */
1302:         {
1303:           DMStagStencil row;

1305:           row.i   = ex;
1306:           row.j   = ey;
1307:           row.k   = ez;
1308:           row.loc = DMSTAG_DOWN;
1309:           row.c   = 0;

1311:           if (bottom_boundary) {
1312:             /* Bottom y-velocity Dirichlet */
1313:             const PetscScalar val_rhs = 0.0;
1314:             const PetscScalar val_A   = K_bound;

1316:             DMStagMatSetValuesStencil(dm_main, A, 1, &row, 1, &row, &val_A, INSERT_VALUES);
1317:             if (build_rhs) DMStagVecSetValuesStencil(dm_main, rhs, 1, &row, &val_rhs, INSERT_VALUES);
1318:           } else {
1319:             /* Y-momentum equation (including non-zero forcing) */
1320:             PetscInt      count;
1321:             DMStagStencil col[17];
1322:             PetscScalar   val_rhs, val_A[17];
1323:             DMStagStencil eta_point[6], rho_point[4];
1324:             PetscScalar   eta[6], rho[4], eta_left, eta_right, eta_up, eta_down, eta_back, eta_front; /* relative to the bottom face */

1326:             if (build_rhs) {
1327:               /* get rho values  (note .c = 1) */
1328:               /* Note that we have rho at perhaps strange points (edges not corners) */
1329:               rho_point[0].i   = ex;
1330:               rho_point[0].j   = ey;
1331:               rho_point[0].k   = ez;
1332:               rho_point[0].loc = DMSTAG_DOWN_LEFT;
1333:               rho_point[0].c   = 1;
1334:               rho_point[1].i   = ex;
1335:               rho_point[1].j   = ey;
1336:               rho_point[1].k   = ez;
1337:               rho_point[1].loc = DMSTAG_DOWN_RIGHT;
1338:               rho_point[1].c   = 1;
1339:               rho_point[2].i   = ex;
1340:               rho_point[2].j   = ey;
1341:               rho_point[2].k   = ez;
1342:               rho_point[2].loc = DMSTAG_BACK_DOWN;
1343:               rho_point[2].c   = 1;
1344:               rho_point[3].i   = ex;
1345:               rho_point[3].j   = ey;
1346:               rho_point[3].k   = ez;
1347:               rho_point[3].loc = DMSTAG_FRONT_DOWN;
1348:               rho_point[3].c   = 1;
1349:               DMStagVecGetValuesStencil(dm_coefficients, coeff_local, 4, rho_point, rho);

1351:               /* Compute forcing */
1352:               val_rhs = ctx->gy * dv * (0.25 * (rho[0] + rho[1] + rho[2] + rho[3]));
1353:             }

1355:             /* Get eta values */
1356:             eta_point[0].i   = ex;
1357:             eta_point[0].j   = ey;
1358:             eta_point[0].k   = ez;
1359:             eta_point[0].loc = DMSTAG_DOWN_LEFT;
1360:             eta_point[0].c   = 0; /* Left  */
1361:             eta_point[1].i   = ex;
1362:             eta_point[1].j   = ey;
1363:             eta_point[1].k   = ez;
1364:             eta_point[1].loc = DMSTAG_DOWN_RIGHT;
1365:             eta_point[1].c   = 0; /* Right */
1366:             eta_point[2].i   = ex;
1367:             eta_point[2].j   = ey - 1;
1368:             eta_point[2].k   = ez;
1369:             eta_point[2].loc = DMSTAG_ELEMENT;
1370:             eta_point[2].c   = 0; /* Down  */
1371:             eta_point[3].i   = ex;
1372:             eta_point[3].j   = ey;
1373:             eta_point[3].k   = ez;
1374:             eta_point[3].loc = DMSTAG_ELEMENT;
1375:             eta_point[3].c   = 0; /* Up    */
1376:             eta_point[4].i   = ex;
1377:             eta_point[4].j   = ey;
1378:             eta_point[4].k   = ez;
1379:             eta_point[4].loc = DMSTAG_BACK_DOWN;
1380:             eta_point[4].c   = 0; /* Back  */
1381:             eta_point[5].i   = ex;
1382:             eta_point[5].j   = ey;
1383:             eta_point[5].k   = ez;
1384:             eta_point[5].loc = DMSTAG_FRONT_DOWN;
1385:             eta_point[5].c   = 0; /* Front  */
1386:             DMStagVecGetValuesStencil(dm_coefficients, coeff_local, 6, eta_point, eta);
1387:             eta_left  = eta[0];
1388:             eta_right = eta[1];
1389:             eta_down  = eta[2];
1390:             eta_up    = eta[3];
1391:             eta_back  = eta[4];
1392:             eta_front = eta[5];

1394:             count = 0;

1396:             col[count]   = row;
1397:             val_A[count] = -2.0 * dv * (eta_up + eta_down) / (hy * hy);
1398:             if (!left_boundary) val_A[count] += -1.0 * dv * eta_left / (hx * hx);
1399:             if (!right_boundary) val_A[count] += -1.0 * dv * eta_right / (hx * hx);
1400:             if (!back_boundary) val_A[count] += -1.0 * dv * eta_back / (hz * hz);
1401:             if (!front_boundary) val_A[count] += -1.0 * dv * eta_front / (hz * hz);
1402:             ++count;

1404:             col[count].i   = ex;
1405:             col[count].j   = ey - 1;
1406:             col[count].k   = ez;
1407:             col[count].loc = DMSTAG_DOWN;
1408:             col[count].c   = 0;
1409:             val_A[count]   = 2.0 * dv * eta_down / (hy * hy);
1410:             ++count;
1411:             col[count].i   = ex;
1412:             col[count].j   = ey + 1;
1413:             col[count].k   = ez;
1414:             col[count].loc = DMSTAG_DOWN;
1415:             col[count].c   = 0;
1416:             val_A[count]   = 2.0 * dv * eta_up / (hy * hy);
1417:             ++count;

1419:             if (!left_boundary) {
1420:               col[count].i   = ex - 1;
1421:               col[count].j   = ey;
1422:               col[count].k   = ez;
1423:               col[count].loc = DMSTAG_DOWN;
1424:               col[count].c   = 0;
1425:               val_A[count]   = dv * eta_left / (hx * hx);
1426:               ++count;
1427:             }
1428:             if (!right_boundary) {
1429:               col[count].i   = ex + 1;
1430:               col[count].j   = ey;
1431:               col[count].k   = ez;
1432:               col[count].loc = DMSTAG_DOWN;
1433:               col[count].c   = 0;
1434:               val_A[count]   = dv * eta_right / (hx * hx);
1435:               ++count;
1436:             }
1437:             if (!back_boundary) {
1438:               col[count].i   = ex;
1439:               col[count].j   = ey;
1440:               col[count].k   = ez - 1;
1441:               col[count].loc = DMSTAG_DOWN;
1442:               col[count].c   = 0;
1443:               val_A[count]   = dv * eta_back / (hz * hz);
1444:               ++count;
1445:             }
1446:             if (!front_boundary) {
1447:               col[count].i   = ex;
1448:               col[count].j   = ey;
1449:               col[count].k   = ez + 1;
1450:               col[count].loc = DMSTAG_DOWN;
1451:               col[count].c   = 0;
1452:               val_A[count]   = dv * eta_front / (hz * hz);
1453:               ++count;
1454:             }

1456:             col[count].i   = ex;
1457:             col[count].j   = ey - 1;
1458:             col[count].k   = ez;
1459:             col[count].loc = DMSTAG_LEFT;
1460:             col[count].c   = 0;
1461:             val_A[count]   = dv * eta_left / (hx * hy);
1462:             ++count; /* down left*/
1463:             col[count].i   = ex;
1464:             col[count].j   = ey;
1465:             col[count].k   = ez;
1466:             col[count].loc = DMSTAG_LEFT;
1467:             col[count].c   = 0;
1468:             val_A[count]   = -1.0 * dv * eta_left / (hx * hy);
1469:             ++count; /* up left*/

1471:             col[count].i   = ex;
1472:             col[count].j   = ey - 1;
1473:             col[count].k   = ez;
1474:             col[count].loc = DMSTAG_RIGHT;
1475:             col[count].c   = 0;
1476:             val_A[count]   = -1.0 * dv * eta_right / (hx * hy);
1477:             ++count; /* down right*/
1478:             col[count].i   = ex;
1479:             col[count].j   = ey;
1480:             col[count].k   = ez;
1481:             col[count].loc = DMSTAG_RIGHT;
1482:             col[count].c   = 0;
1483:             val_A[count]   = dv * eta_right / (hx * hy);
1484:             ++count; /* up right*/

1486:             col[count].i   = ex;
1487:             col[count].j   = ey - 1;
1488:             col[count].k   = ez;
1489:             col[count].loc = DMSTAG_BACK;
1490:             col[count].c   = 0;
1491:             val_A[count]   = dv * eta_back / (hy * hz);
1492:             ++count; /* back down */
1493:             col[count].i   = ex;
1494:             col[count].j   = ey;
1495:             col[count].k   = ez;
1496:             col[count].loc = DMSTAG_BACK;
1497:             col[count].c   = 0;
1498:             val_A[count]   = -1.0 * dv * eta_back / (hy * hz);
1499:             ++count; /* back up */

1501:             col[count].i   = ex;
1502:             col[count].j   = ey - 1;
1503:             col[count].k   = ez;
1504:             col[count].loc = DMSTAG_FRONT;
1505:             col[count].c   = 0;
1506:             val_A[count]   = -1.0 * dv * eta_front / (hy * hz);
1507:             ++count; /* front down */
1508:             col[count].i   = ex;
1509:             col[count].j   = ey;
1510:             col[count].k   = ez;
1511:             col[count].loc = DMSTAG_FRONT;
1512:             col[count].c   = 0;
1513:             val_A[count]   = dv * eta_front / (hy * hz);
1514:             ++count; /* front up */

1516:             if (!parameters->faces_only) {
1517:               col[count].i   = ex;
1518:               col[count].j   = ey - 1;
1519:               col[count].k   = ez;
1520:               col[count].loc = DMSTAG_ELEMENT;
1521:               col[count].c   = 0;
1522:               val_A[count]   = K_cont * dv / hy;
1523:               ++count;
1524:               col[count].i   = ex;
1525:               col[count].j   = ey;
1526:               col[count].k   = ez;
1527:               col[count].loc = DMSTAG_ELEMENT;
1528:               col[count].c   = 0;
1529:               val_A[count]   = -1.0 * K_cont * dv / hy;
1530:               ++count;
1531:             }

1533:             DMStagMatSetValuesStencil(dm_main, A, 1, &row, count, col, val_A, INSERT_VALUES);
1534:             if (build_rhs) DMStagVecSetValuesStencil(dm_main, rhs, 1, &row, &val_rhs, INSERT_VALUES);
1535:           }
1536:         }

1538:         if (front_boundary) {
1539:           /* Front z-velocity Dirichlet */
1540:           DMStagStencil     row;
1541:           const PetscScalar val_rhs = 0.0;
1542:           const PetscScalar val_A   = K_bound;

1544:           row.i   = ex;
1545:           row.j   = ey;
1546:           row.k   = ez;
1547:           row.loc = DMSTAG_FRONT;
1548:           row.c   = 0;
1549:           DMStagMatSetValuesStencil(dm_main, A, 1, &row, 1, &row, &val_A, INSERT_VALUES);
1550:           if (build_rhs) DMStagVecSetValuesStencil(dm_main, rhs, 1, &row, &val_rhs, INSERT_VALUES);
1551:         }

1553:         /* Z faces - back */
1554:         {
1555:           DMStagStencil row;

1557:           row.i   = ex;
1558:           row.j   = ey;
1559:           row.k   = ez;
1560:           row.loc = DMSTAG_BACK;
1561:           row.c   = 0;

1563:           if (back_boundary) {
1564:             /* Back z-velocity Dirichlet */
1565:             const PetscScalar val_rhs = 0.0;
1566:             const PetscScalar val_A   = K_bound;

1568:             DMStagMatSetValuesStencil(dm_main, A, 1, &row, 1, &row, &val_A, INSERT_VALUES);
1569:             if (build_rhs) DMStagVecSetValuesStencil(dm_main, rhs, 1, &row, &val_rhs, INSERT_VALUES);
1570:           } else {
1571:             /* Z-momentum equation */
1572:             PetscInt          count;
1573:             DMStagStencil     col[17];
1574:             PetscScalar       val_A[17];
1575:             DMStagStencil     eta_point[6];
1576:             PetscScalar       eta[6], eta_left, eta_right, eta_up, eta_down, eta_back, eta_front; /* relative to the back face */
1577:             const PetscScalar val_rhs = 0.0;

1579:             /* Get eta values */
1580:             eta_point[0].i   = ex;
1581:             eta_point[0].j   = ey;
1582:             eta_point[0].k   = ez;
1583:             eta_point[0].loc = DMSTAG_BACK_LEFT;
1584:             eta_point[0].c   = 0; /* Left  */
1585:             eta_point[1].i   = ex;
1586:             eta_point[1].j   = ey;
1587:             eta_point[1].k   = ez;
1588:             eta_point[1].loc = DMSTAG_BACK_RIGHT;
1589:             eta_point[1].c   = 0; /* Right */
1590:             eta_point[2].i   = ex;
1591:             eta_point[2].j   = ey;
1592:             eta_point[2].k   = ez;
1593:             eta_point[2].loc = DMSTAG_BACK_DOWN;
1594:             eta_point[2].c   = 0; /* Down  */
1595:             eta_point[3].i   = ex;
1596:             eta_point[3].j   = ey;
1597:             eta_point[3].k   = ez;
1598:             eta_point[3].loc = DMSTAG_BACK_UP;
1599:             eta_point[3].c   = 0; /* Up    */
1600:             eta_point[4].i   = ex;
1601:             eta_point[4].j   = ey;
1602:             eta_point[4].k   = ez - 1;
1603:             eta_point[4].loc = DMSTAG_ELEMENT;
1604:             eta_point[4].c   = 0; /* Back  */
1605:             eta_point[5].i   = ex;
1606:             eta_point[5].j   = ey;
1607:             eta_point[5].k   = ez;
1608:             eta_point[5].loc = DMSTAG_ELEMENT;
1609:             eta_point[5].c   = 0; /* Front  */
1610:             DMStagVecGetValuesStencil(dm_coefficients, coeff_local, 6, eta_point, eta);
1611:             eta_left  = eta[0];
1612:             eta_right = eta[1];
1613:             eta_down  = eta[2];
1614:             eta_up    = eta[3];
1615:             eta_back  = eta[4];
1616:             eta_front = eta[5];

1618:             count = 0;

1620:             col[count]   = row;
1621:             val_A[count] = -2.0 * dv * (eta_back + eta_front) / (hz * hz);
1622:             if (!left_boundary) val_A[count] += -1.0 * dv * eta_left / (hx * hx);
1623:             if (!right_boundary) val_A[count] += -1.0 * dv * eta_right / (hx * hx);
1624:             if (!top_boundary) val_A[count] += -1.0 * dv * eta_up / (hy * hy);
1625:             if (!bottom_boundary) val_A[count] += -1.0 * dv * eta_down / (hy * hy);
1626:             ++count;

1628:             col[count].i   = ex;
1629:             col[count].j   = ey;
1630:             col[count].k   = ez - 1;
1631:             col[count].loc = DMSTAG_BACK;
1632:             col[count].c   = 0;
1633:             val_A[count]   = 2.0 * dv * eta_back / (hz * hz);
1634:             ++count;
1635:             col[count].i   = ex;
1636:             col[count].j   = ey;
1637:             col[count].k   = ez + 1;
1638:             col[count].loc = DMSTAG_BACK;
1639:             col[count].c   = 0;
1640:             val_A[count]   = 2.0 * dv * eta_front / (hz * hz);
1641:             ++count;

1643:             if (!left_boundary) {
1644:               col[count].i   = ex - 1;
1645:               col[count].j   = ey;
1646:               col[count].k   = ez;
1647:               col[count].loc = DMSTAG_BACK;
1648:               col[count].c   = 0;
1649:               val_A[count]   = dv * eta_left / (hx * hx);
1650:               ++count;
1651:             }
1652:             if (!right_boundary) {
1653:               col[count].i   = ex + 1;
1654:               col[count].j   = ey;
1655:               col[count].k   = ez;
1656:               col[count].loc = DMSTAG_BACK;
1657:               col[count].c   = 0;
1658:               val_A[count]   = dv * eta_right / (hx * hx);
1659:               ++count;
1660:             }
1661:             if (!bottom_boundary) {
1662:               col[count].i   = ex;
1663:               col[count].j   = ey - 1;
1664:               col[count].k   = ez;
1665:               col[count].loc = DMSTAG_BACK;
1666:               col[count].c   = 0;
1667:               val_A[count]   = dv * eta_down / (hy * hy);
1668:               ++count;
1669:             }
1670:             if (!top_boundary) {
1671:               col[count].i   = ex;
1672:               col[count].j   = ey + 1;
1673:               col[count].k   = ez;
1674:               col[count].loc = DMSTAG_BACK;
1675:               col[count].c   = 0;
1676:               val_A[count]   = dv * eta_up / (hy * hy);
1677:               ++count;
1678:             }

1680:             col[count].i   = ex;
1681:             col[count].j   = ey;
1682:             col[count].k   = ez - 1;
1683:             col[count].loc = DMSTAG_LEFT;
1684:             col[count].c   = 0;
1685:             val_A[count]   = dv * eta_left / (hx * hz);
1686:             ++count; /* back left*/
1687:             col[count].i   = ex;
1688:             col[count].j   = ey;
1689:             col[count].k   = ez;
1690:             col[count].loc = DMSTAG_LEFT;
1691:             col[count].c   = 0;
1692:             val_A[count]   = -1.0 * dv * eta_left / (hx * hz);
1693:             ++count; /* front left*/

1695:             col[count].i   = ex;
1696:             col[count].j   = ey;
1697:             col[count].k   = ez - 1;
1698:             col[count].loc = DMSTAG_RIGHT;
1699:             col[count].c   = 0;
1700:             val_A[count]   = -1.0 * dv * eta_right / (hx * hz);
1701:             ++count; /* back right */
1702:             col[count].i   = ex;
1703:             col[count].j   = ey;
1704:             col[count].k   = ez;
1705:             col[count].loc = DMSTAG_RIGHT;
1706:             col[count].c   = 0;
1707:             val_A[count]   = dv * eta_right / (hx * hz);
1708:             ++count; /* front right*/

1710:             col[count].i   = ex;
1711:             col[count].j   = ey;
1712:             col[count].k   = ez - 1;
1713:             col[count].loc = DMSTAG_DOWN;
1714:             col[count].c   = 0;
1715:             val_A[count]   = dv * eta_down / (hy * hz);
1716:             ++count; /* back down */
1717:             col[count].i   = ex;
1718:             col[count].j   = ey;
1719:             col[count].k   = ez;
1720:             col[count].loc = DMSTAG_DOWN;
1721:             col[count].c   = 0;
1722:             val_A[count]   = -1.0 * dv * eta_down / (hy * hz);
1723:             ++count; /* back down */

1725:             col[count].i   = ex;
1726:             col[count].j   = ey;
1727:             col[count].k   = ez - 1;
1728:             col[count].loc = DMSTAG_UP;
1729:             col[count].c   = 0;
1730:             val_A[count]   = -1.0 * dv * eta_up / (hy * hz);
1731:             ++count; /* back up */
1732:             col[count].i   = ex;
1733:             col[count].j   = ey;
1734:             col[count].k   = ez;
1735:             col[count].loc = DMSTAG_UP;
1736:             col[count].c   = 0;
1737:             val_A[count]   = dv * eta_up / (hy * hz);
1738:             ++count; /* back up */

1740:             if (!parameters->faces_only) {
1741:               col[count].i   = ex;
1742:               col[count].j   = ey;
1743:               col[count].k   = ez - 1;
1744:               col[count].loc = DMSTAG_ELEMENT;
1745:               col[count].c   = 0;
1746:               val_A[count]   = K_cont * dv / hz;
1747:               ++count;
1748:               col[count].i   = ex;
1749:               col[count].j   = ey;
1750:               col[count].k   = ez;
1751:               col[count].loc = DMSTAG_ELEMENT;
1752:               col[count].c   = 0;
1753:               val_A[count]   = -1.0 * K_cont * dv / hz;
1754:               ++count;
1755:             }

1757:             DMStagMatSetValuesStencil(dm_main, A, 1, &row, count, col, val_A, INSERT_VALUES);
1758:             if (build_rhs) DMStagVecSetValuesStencil(dm_main, rhs, 1, &row, &val_rhs, INSERT_VALUES);
1759:           }
1760:         }

1762:         /* Elements */
1763:         if (!parameters->faces_only) {
1764:           DMStagStencil row;

1766:           row.i   = ex;
1767:           row.j   = ey;
1768:           row.k   = ez;
1769:           row.loc = DMSTAG_ELEMENT;
1770:           row.c   = 0;

1772:           if (ctx->pin_pressure && ex == pinx && ey == piny && ez == pinz) {
1773:             /* Pin a pressure node to zero, if requested */
1774:             const PetscScalar val_A   = K_bound;
1775:             const PetscScalar val_rhs = 0.0;

1777:             DMStagMatSetValuesStencil(dm_main, A, 1, &row, 1, &row, &val_A, INSERT_VALUES);
1778:             if (build_rhs) DMStagVecSetValuesStencil(dm_main, rhs, 1, &row, &val_rhs, INSERT_VALUES);
1779:           } else {
1780:             /* Continuity equation */
1781:             /* Note that this includes an explicit zero on the diagonal. This is only needed for
1782:                some direct solvers (not required if using an iterative solver and setting a constant-pressure nullspace) */
1783:             /* Note: the scaling by dv is not chosen in a principled way and is likely sub-optimal */
1784:             DMStagStencil     col[7];
1785:             PetscScalar       val_A[7];
1786:             const PetscScalar val_rhs = 0.0;

1788:             col[0].i   = ex;
1789:             col[0].j   = ey;
1790:             col[0].k   = ez;
1791:             col[0].loc = DMSTAG_LEFT;
1792:             col[0].c   = 0;
1793:             val_A[0]   = -1.0 * K_cont * dv / hx;
1794:             col[1].i   = ex;
1795:             col[1].j   = ey;
1796:             col[1].k   = ez;
1797:             col[1].loc = DMSTAG_RIGHT;
1798:             col[1].c   = 0;
1799:             val_A[1]   = K_cont * dv / hx;
1800:             col[2].i   = ex;
1801:             col[2].j   = ey;
1802:             col[2].k   = ez;
1803:             col[2].loc = DMSTAG_DOWN;
1804:             col[2].c   = 0;
1805:             val_A[2]   = -1.0 * K_cont * dv / hy;
1806:             col[3].i   = ex;
1807:             col[3].j   = ey;
1808:             col[3].k   = ez;
1809:             col[3].loc = DMSTAG_UP;
1810:             col[3].c   = 0;
1811:             val_A[3]   = K_cont * dv / hy;
1812:             col[4].i   = ex;
1813:             col[4].j   = ey;
1814:             col[4].k   = ez;
1815:             col[4].loc = DMSTAG_BACK;
1816:             col[4].c   = 0;
1817:             val_A[4]   = -1.0 * K_cont * dv / hz;
1818:             col[5].i   = ex;
1819:             col[5].j   = ey;
1820:             col[5].k   = ez;
1821:             col[5].loc = DMSTAG_FRONT;
1822:             col[5].c   = 0;
1823:             val_A[5]   = K_cont * dv / hz;
1824:             col[6]     = row;
1825:             val_A[6]   = 0.0;
1826:             DMStagMatSetValuesStencil(dm_main, A, 1, &row, 7, col, val_A, INSERT_VALUES);
1827:             if (build_rhs) DMStagVecSetValuesStencil(dm_main, rhs, 1, &row, &val_rhs, INSERT_VALUES);
1828:           }
1829:         }
1830:       }
1831:     }
1832:   }
1833:   DMRestoreLocalVector(dm_coefficients, &coeff_local);

1835:   /* Add additional inverse viscosity terms (for use in building a preconditioning matrix) */
1836:   if (parameters->include_inverse_visc) {
1838:     OperatorInsertInverseViscosityPressureTerms(dm_main, dm_coefficients, ctx->levels[level]->coeff, 1.0, A);
1839:   }

1841:   MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY);
1842:   if (build_rhs) VecAssemblyBegin(rhs);
1843:   MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY);
1844:   if (build_rhs) VecAssemblyEnd(rhs);
1845:   return 0;
1846: }

1848: static PetscErrorCode CreateSystem(SystemParameters parameters, Mat *pA, Vec *pRhs)
1849: {
1851:   if (parameters->ctx->dim == 2) {
1852:     CreateSystem2d(parameters, pA, pRhs);
1853:   } else if (parameters->ctx->dim == 3) {
1854:     CreateSystem3d(parameters, pA, pRhs);
1855:   } else SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_SUP, "Unsupported dimension %" PetscInt_FMT, parameters->ctx->dim);
1856:   return 0;
1857: }

1859: PetscErrorCode PopulateCoefficientData(Ctx ctx, PetscInt level)
1860: {
1861:   PetscInt    dim;
1862:   PetscInt    N[3];
1863:   PetscInt    ex, ey, ez, startx, starty, startz, nx, ny, nz;
1864:   PetscInt    slot_prev, slot_center;
1865:   PetscInt    slot_rho_downleft, slot_rho_backleft, slot_rho_backdown, slot_eta_element, slot_eta_downleft, slot_eta_backleft, slot_eta_backdown;
1866:   Vec         coeff_local;
1867:   PetscReal **arr_coordinates_x, **arr_coordinates_y, **arr_coordinates_z;
1868:   DM          dm_coefficients;
1869:   Vec         coeff;

1872:   dm_coefficients = ctx->levels[level]->dm_coefficients;
1873:   DMGetDimension(dm_coefficients, &dim);

1875:   /* Create global coefficient vector */
1876:   DMCreateGlobalVector(dm_coefficients, &ctx->levels[level]->coeff);
1877:   coeff = ctx->levels[level]->coeff;

1879:   /* Get temporary access to a local representation of the coefficient data */
1880:   DMGetLocalVector(dm_coefficients, &coeff_local);
1881:   DMGlobalToLocal(dm_coefficients, coeff, INSERT_VALUES, coeff_local);

1883:   /* Use direct array access to coefficient and coordinate arrays, to popoulate coefficient data */
1884:   DMStagGetGhostCorners(dm_coefficients, &startx, &starty, &startz, &nx, &ny, &nz);
1885:   DMStagGetGlobalSizes(dm_coefficients, &N[0], &N[1], &N[2]);
1886:   DMStagGetProductCoordinateArraysRead(dm_coefficients, &arr_coordinates_x, &arr_coordinates_y, &arr_coordinates_z);
1887:   DMStagGetProductCoordinateLocationSlot(dm_coefficients, DMSTAG_ELEMENT, &slot_center);
1888:   DMStagGetProductCoordinateLocationSlot(dm_coefficients, DMSTAG_LEFT, &slot_prev);
1889:   DMStagGetLocationSlot(dm_coefficients, DMSTAG_ELEMENT, 0, &slot_eta_element);
1890:   DMStagGetLocationSlot(dm_coefficients, DMSTAG_DOWN_LEFT, 0, &slot_eta_downleft);
1891:   DMStagGetLocationSlot(dm_coefficients, DMSTAG_DOWN_LEFT, 1, &slot_rho_downleft);
1892:   if (dim == 2) {
1893:     PetscScalar ***arr_coefficients;

1895:     DMStagVecGetArray(dm_coefficients, coeff_local, &arr_coefficients);
1896:     /* Note that these ranges are with respect to the local representation */
1897:     for (ey = starty; ey < starty + ny; ++ey) {
1898:       for (ex = startx; ex < startx + nx; ++ex) {
1899:         arr_coefficients[ey][ex][slot_eta_element]  = ctx->GetEta(ctx, arr_coordinates_x[ex][slot_center], arr_coordinates_y[ey][slot_center], 0.0);
1900:         arr_coefficients[ey][ex][slot_eta_downleft] = ctx->GetEta(ctx, arr_coordinates_x[ex][slot_prev], arr_coordinates_y[ey][slot_prev], 0.0);
1901:         arr_coefficients[ey][ex][slot_rho_downleft] = ctx->GetRho(ctx, arr_coordinates_x[ex][slot_prev], arr_coordinates_y[ey][slot_prev], 0.0);
1902:       }
1903:     }
1904:     DMStagVecRestoreArray(dm_coefficients, coeff_local, &arr_coefficients);
1905:   } else if (dim == 3) {
1906:     PetscScalar ****arr_coefficients;

1908:     DMStagGetLocationSlot(dm_coefficients, DMSTAG_BACK_LEFT, 0, &slot_eta_backleft);
1909:     DMStagGetLocationSlot(dm_coefficients, DMSTAG_BACK_LEFT, 1, &slot_rho_backleft);
1910:     DMStagGetLocationSlot(dm_coefficients, DMSTAG_BACK_DOWN, 0, &slot_eta_backdown);
1911:     DMStagGetLocationSlot(dm_coefficients, DMSTAG_BACK_DOWN, 1, &slot_rho_backdown);
1912:     DMStagVecGetArray(dm_coefficients, coeff_local, &arr_coefficients);
1913:     /* Note that these are with respect to the entire local representation, including ghosts */
1914:     for (ez = startz; ez < startz + nz; ++ez) {
1915:       for (ey = starty; ey < starty + ny; ++ey) {
1916:         for (ex = startx; ex < startx + nx; ++ex) {
1917:           const PetscScalar x_prev   = arr_coordinates_x[ex][slot_prev];
1918:           const PetscScalar y_prev   = arr_coordinates_y[ey][slot_prev];
1919:           const PetscScalar z_prev   = arr_coordinates_z[ez][slot_prev];
1920:           const PetscScalar x_center = arr_coordinates_x[ex][slot_center];
1921:           const PetscScalar y_center = arr_coordinates_y[ey][slot_center];
1922:           const PetscScalar z_center = arr_coordinates_z[ez][slot_center];

1924:           arr_coefficients[ez][ey][ex][slot_eta_element]  = ctx->GetEta(ctx, x_center, y_center, z_center);
1925:           arr_coefficients[ez][ey][ex][slot_eta_downleft] = ctx->GetEta(ctx, x_prev, y_prev, z_center);
1926:           arr_coefficients[ez][ey][ex][slot_rho_downleft] = ctx->GetRho(ctx, x_prev, y_prev, z_center);
1927:           arr_coefficients[ez][ey][ex][slot_eta_backleft] = ctx->GetEta(ctx, x_prev, y_center, z_prev);
1928:           arr_coefficients[ez][ey][ex][slot_rho_backleft] = ctx->GetRho(ctx, x_prev, y_center, z_prev);
1929:           arr_coefficients[ez][ey][ex][slot_eta_backdown] = ctx->GetEta(ctx, x_center, y_prev, z_prev);
1930:           arr_coefficients[ez][ey][ex][slot_rho_backdown] = ctx->GetRho(ctx, x_center, y_prev, z_prev);
1931:         }
1932:       }
1933:     }
1934:     DMStagVecRestoreArray(dm_coefficients, coeff_local, &arr_coefficients);
1935:   } else SETERRQ(PetscObjectComm((PetscObject)dm_coefficients), PETSC_ERR_SUP, "Unsupported dimension %" PetscInt_FMT, dim);
1936:   DMStagRestoreProductCoordinateArraysRead(dm_coefficients, &arr_coordinates_x, &arr_coordinates_y, &arr_coordinates_z);
1937:   DMLocalToGlobal(dm_coefficients, coeff_local, INSERT_VALUES, coeff);
1938:   DMRestoreLocalVector(dm_coefficients, &coeff_local);
1939:   return 0;
1940: }

1942: static PetscErrorCode CreateAuxiliaryOperator(Ctx ctx, PetscInt level, Mat *p_S_hat)
1943: {
1944:   DM  dm_element;
1945:   Mat S_hat;
1946:   DM  dm_stokes, dm_coefficients;
1947:   Vec coeff;

1950:   dm_stokes       = ctx->levels[level]->dm_stokes;
1951:   dm_coefficients = ctx->levels[level]->dm_coefficients;
1952:   coeff           = ctx->levels[level]->coeff;
1953:   if (ctx->dim == 2) {
1954:     DMStagCreateCompatibleDMStag(dm_stokes, 0, 0, 1, 0, &dm_element);
1955:   } else if (ctx->dim == 3) {
1956:     DMStagCreateCompatibleDMStag(dm_stokes, 0, 0, 0, 1, &dm_element);
1957:   } else SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_SUP, "Not implemented for dimension %" PetscInt_FMT, ctx->dim);
1958:   DMCreateMatrix(dm_element, p_S_hat);
1959:   S_hat = *p_S_hat;
1960:   OperatorInsertInverseViscosityPressureTerms(dm_element, dm_coefficients, coeff, 1.0, S_hat);
1961:   MatAssemblyBegin(S_hat, MAT_FINAL_ASSEMBLY);
1962:   MatAssemblyEnd(S_hat, MAT_FINAL_ASSEMBLY);
1963:   DMDestroy(&dm_element);
1964:   return 0;
1965: }

1967: static PetscErrorCode OperatorInsertInverseViscosityPressureTerms(DM dm, DM dm_coefficients, Vec coefficients, PetscScalar scale, Mat mat)
1968: {
1969:   PetscInt dim, ex, ey, ez, startx, starty, startz, nx, ny, nz;
1970:   Vec      coeff_local;

1973:   DMGetDimension(dm, &dim);
1974:   DMGetLocalVector(dm_coefficients, &coeff_local);
1975:   DMGlobalToLocal(dm_coefficients, coefficients, INSERT_VALUES, coeff_local);
1976:   DMStagGetCorners(dm, &startx, &starty, &startz, &nx, &ny, &nz, NULL, NULL, NULL);
1977:   if (dim == 2) { /* Trick to have one loop nest */
1978:     startz = 0;
1979:     nz     = 1;
1980:   }
1981:   for (ez = startz; ez < startz + nz; ++ez) {
1982:     for (ey = starty; ey < starty + ny; ++ey) {
1983:       for (ex = startx; ex < startx + nx; ++ex) {
1984:         DMStagStencil from, to;
1985:         PetscScalar   val;

1987:         /* component 0 on element is viscosity */
1988:         from.i   = ex;
1989:         from.j   = ey;
1990:         from.k   = ez;
1991:         from.c   = 0;
1992:         from.loc = DMSTAG_ELEMENT;
1993:         DMStagVecGetValuesStencil(dm_coefficients, coeff_local, 1, &from, &val);
1994:         val = scale / val; /* inverse viscosity, scaled */
1995:         to  = from;
1996:         DMStagMatSetValuesStencil(dm, mat, 1, &to, 1, &to, &val, INSERT_VALUES);
1997:       }
1998:     }
1999:   }
2000:   DMRestoreLocalVector(dm_coefficients, &coeff_local);
2001:   /* Note that this function does not call MatAssembly{Begin,End} */
2002:   return 0;
2003: }

2005: /* Create a pressure-only DMStag and use it to generate a nullspace vector
2006:    - Create a compatible DMStag with one dof per element (and nothing else).
2007:    - Create a constant vector and normalize it
2008:    - Migrate it to a vector on the original dmSol (making use of the fact
2009:    that this will fill in zeros for "extra" dof)
2010:    - Set the nullspace for the operator
2011:    - Destroy everything (the operator keeps the references it needs) */
2012: static PetscErrorCode AttachNullspace(DM dmSol, Mat A)
2013: {
2014:   DM           dmPressure;
2015:   Vec          constantPressure, basis;
2016:   PetscReal    nrm;
2017:   MatNullSpace matNullSpace;

2020:   DMStagCreateCompatibleDMStag(dmSol, 0, 0, 1, 0, &dmPressure);
2021:   DMGetGlobalVector(dmPressure, &constantPressure);
2022:   VecSet(constantPressure, 1.0);
2023:   VecNorm(constantPressure, NORM_2, &nrm);
2024:   VecScale(constantPressure, 1.0 / nrm);
2025:   DMCreateGlobalVector(dmSol, &basis);
2026:   DMStagMigrateVec(dmPressure, constantPressure, dmSol, basis);
2027:   MatNullSpaceCreate(PetscObjectComm((PetscObject)dmSol), PETSC_FALSE, 1, &basis, &matNullSpace);
2028:   VecDestroy(&basis);
2029:   MatSetNullSpace(A, matNullSpace);
2030:   MatNullSpaceDestroy(&matNullSpace);
2031:   DMRestoreGlobalVector(dmPressure, &constantPressure);
2032:   DMDestroy(&dmPressure);
2033:   return 0;
2034: }

2036: static PetscErrorCode DumpSolution(Ctx ctx, PetscInt level, Vec x)
2037: {
2038:   DM       dm_stokes, dm_coefficients;
2039:   Vec      coeff;
2040:   DM       dm_vel_avg;
2041:   Vec      vel_avg;
2042:   DM       da_vel_avg, da_p, da_eta_element;
2043:   Vec      vec_vel_avg, vec_p, vec_eta_element;
2044:   DM       da_eta_down_left, da_rho_down_left, da_eta_back_left, da_rho_back_left, da_eta_back_down, da_rho_back_down;
2045:   Vec      vec_eta_down_left, vec_rho_down_left, vec_eta_back_left, vec_rho_back_left, vec_eta_back_down, vec_rho_back_down;
2046:   PetscInt ex, ey, ez, startx, starty, startz, nx, ny, nz;
2047:   Vec      stokesLocal;

2050:   dm_stokes       = ctx->levels[level]->dm_stokes;
2051:   dm_coefficients = ctx->levels[level]->dm_coefficients;
2052:   coeff           = ctx->levels[level]->coeff;

2054:   /* For convenience, create a new DM and Vec which will hold averaged velocities
2055:      Note that this could also be accomplished with direct array access, using
2056:      DMStagVecGetArray() and related functions */
2057:   if (ctx->dim == 2) {
2058:     DMStagCreateCompatibleDMStag(dm_stokes, 0, 0, 2, 0, &dm_vel_avg); /* 2 dof per element */
2059:   } else if (ctx->dim == 3) {
2060:     DMStagCreateCompatibleDMStag(dm_stokes, 0, 0, 0, 3, &dm_vel_avg); /* 3 dof per element */
2061:   } else SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_SUP, "Not Implemented for dimension %" PetscInt_FMT, ctx->dim);
2062:   DMSetUp(dm_vel_avg);
2063:   DMStagSetUniformCoordinatesProduct(dm_vel_avg, ctx->xmin, ctx->xmax, ctx->ymin, ctx->ymax, ctx->zmin, ctx->zmax);
2064:   DMCreateGlobalVector(dm_vel_avg, &vel_avg);
2065:   DMGetLocalVector(dm_stokes, &stokesLocal);
2066:   DMGlobalToLocal(dm_stokes, x, INSERT_VALUES, stokesLocal);
2067:   DMStagGetCorners(dm_vel_avg, &startx, &starty, &startz, &nx, &ny, &nz, NULL, NULL, NULL);
2068:   if (ctx->dim == 2) {
2069:     for (ey = starty; ey < starty + ny; ++ey) {
2070:       for (ex = startx; ex < startx + nx; ++ex) {
2071:         DMStagStencil from[4], to[2];
2072:         PetscScalar   valFrom[4], valTo[2];

2074:         from[0].i   = ex;
2075:         from[0].j   = ey;
2076:         from[0].loc = DMSTAG_UP;
2077:         from[0].c   = 0;
2078:         from[1].i   = ex;
2079:         from[1].j   = ey;
2080:         from[1].loc = DMSTAG_DOWN;
2081:         from[1].c   = 0;
2082:         from[2].i   = ex;
2083:         from[2].j   = ey;
2084:         from[2].loc = DMSTAG_LEFT;
2085:         from[2].c   = 0;
2086:         from[3].i   = ex;
2087:         from[3].j   = ey;
2088:         from[3].loc = DMSTAG_RIGHT;
2089:         from[3].c   = 0;
2090:         DMStagVecGetValuesStencil(dm_stokes, stokesLocal, 4, from, valFrom);
2091:         to[0].i   = ex;
2092:         to[0].j   = ey;
2093:         to[0].loc = DMSTAG_ELEMENT;
2094:         to[0].c   = 0;
2095:         valTo[0]  = 0.5 * (valFrom[2] + valFrom[3]);
2096:         to[1].i   = ex;
2097:         to[1].j   = ey;
2098:         to[1].loc = DMSTAG_ELEMENT;
2099:         to[1].c   = 1;
2100:         valTo[1]  = 0.5 * (valFrom[0] + valFrom[1]);
2101:         DMStagVecSetValuesStencil(dm_vel_avg, vel_avg, 2, to, valTo, INSERT_VALUES);
2102:       }
2103:     }
2104:   } else if (ctx->dim == 3) {
2105:     for (ez = startz; ez < startz + nz; ++ez) {
2106:       for (ey = starty; ey < starty + ny; ++ey) {
2107:         for (ex = startx; ex < startx + nx; ++ex) {
2108:           DMStagStencil from[6], to[3];
2109:           PetscScalar   valFrom[6], valTo[3];

2111:           from[0].i   = ex;
2112:           from[0].j   = ey;
2113:           from[0].k   = ez;
2114:           from[0].loc = DMSTAG_UP;
2115:           from[0].c   = 0;
2116:           from[1].i   = ex;
2117:           from[1].j   = ey;
2118:           from[1].k   = ez;
2119:           from[1].loc = DMSTAG_DOWN;
2120:           from[1].c   = 0;
2121:           from[2].i   = ex;
2122:           from[2].j   = ey;
2123:           from[2].k   = ez;
2124:           from[2].loc = DMSTAG_LEFT;
2125:           from[2].c   = 0;
2126:           from[3].i   = ex;
2127:           from[3].j   = ey;
2128:           from[3].k   = ez;
2129:           from[3].loc = DMSTAG_RIGHT;
2130:           from[3].c   = 0;
2131:           from[4].i   = ex;
2132:           from[4].j   = ey;
2133:           from[4].k   = ez;
2134:           from[4].loc = DMSTAG_BACK;
2135:           from[4].c   = 0;
2136:           from[5].i   = ex;
2137:           from[5].j   = ey;
2138:           from[5].k   = ez;
2139:           from[5].loc = DMSTAG_FRONT;
2140:           from[5].c   = 0;
2141:           DMStagVecGetValuesStencil(dm_stokes, stokesLocal, 6, from, valFrom);
2142:           to[0].i   = ex;
2143:           to[0].j   = ey;
2144:           to[0].k   = ez;
2145:           to[0].loc = DMSTAG_ELEMENT;
2146:           to[0].c   = 0;
2147:           valTo[0]  = 0.5 * (valFrom[2] + valFrom[3]);
2148:           to[1].i   = ex;
2149:           to[1].j   = ey;
2150:           to[1].k   = ez;
2151:           to[1].loc = DMSTAG_ELEMENT;
2152:           to[1].c   = 1;
2153:           valTo[1]  = 0.5 * (valFrom[0] + valFrom[1]);
2154:           to[2].i   = ex;
2155:           to[2].j   = ey;
2156:           to[2].k   = ez;
2157:           to[2].loc = DMSTAG_ELEMENT;
2158:           to[2].c   = 2;
2159:           valTo[2]  = 0.5 * (valFrom[4] + valFrom[5]);
2160:           DMStagVecSetValuesStencil(dm_vel_avg, vel_avg, 3, to, valTo, INSERT_VALUES);
2161:         }
2162:       }
2163:     }
2164:   }
2165:   VecAssemblyBegin(vel_avg);
2166:   VecAssemblyEnd(vel_avg);
2167:   DMRestoreLocalVector(dm_stokes, &stokesLocal);

2169:   /* Create individual DMDAs for sub-grids of our DMStag objects. This is
2170:      somewhat inefficient, but allows use of the DMDA API without re-implementing
2171:      all utilities for DMStag */

2173:   DMStagVecSplitToDMDA(dm_stokes, x, DMSTAG_ELEMENT, 0, &da_p, &vec_p);
2174:   PetscObjectSetName((PetscObject)vec_p, "p (scaled)");

2176:   DMStagVecSplitToDMDA(dm_coefficients, coeff, DMSTAG_ELEMENT, 0, &da_eta_element, &vec_eta_element);
2177:   PetscObjectSetName((PetscObject)vec_eta_element, "eta");

2179:   DMStagVecSplitToDMDA(dm_coefficients, coeff, DMSTAG_DOWN_LEFT, 0, &da_eta_down_left, &vec_eta_down_left);
2180:   PetscObjectSetName((PetscObject)vec_eta_down_left, "eta");

2182:   DMStagVecSplitToDMDA(dm_coefficients, coeff, DMSTAG_DOWN_LEFT, 1, &da_rho_down_left, &vec_rho_down_left);
2183:   PetscObjectSetName((PetscObject)vec_rho_down_left, "density");

2185:   if (ctx->dim == 3) {
2186:     DMStagVecSplitToDMDA(dm_coefficients, coeff, DMSTAG_BACK_LEFT, 0, &da_eta_back_left, &vec_eta_back_left);
2187:     PetscObjectSetName((PetscObject)vec_eta_back_left, "eta");

2189:     DMStagVecSplitToDMDA(dm_coefficients, coeff, DMSTAG_BACK_LEFT, 1, &da_rho_back_left, &vec_rho_back_left);
2190:     PetscObjectSetName((PetscObject)vec_rho_back_left, "rho");

2192:     DMStagVecSplitToDMDA(dm_coefficients, coeff, DMSTAG_BACK_DOWN, 0, &da_eta_back_down, &vec_eta_back_down);
2193:     PetscObjectSetName((PetscObject)vec_eta_back_down, "eta");

2195:     DMStagVecSplitToDMDA(dm_coefficients, coeff, DMSTAG_BACK_DOWN, 1, &da_rho_back_down, &vec_rho_back_down);
2196:     PetscObjectSetName((PetscObject)vec_rho_back_down, "rho");
2197:   }

2199:   DMStagVecSplitToDMDA(dm_vel_avg, vel_avg, DMSTAG_ELEMENT, -3, &da_vel_avg, &vec_vel_avg); /* note -3 : pad with zero */
2200:   PetscObjectSetName((PetscObject)vec_vel_avg, "Velocity (Averaged)");

2202:   /* Dump element-based fields to a .vtr file */
2203:   {
2204:     PetscViewer viewer;

2206:     PetscViewerVTKOpen(PetscObjectComm((PetscObject)da_vel_avg), "ex4_element.vtr", FILE_MODE_WRITE, &viewer);
2207:     VecView(vec_vel_avg, viewer);
2208:     VecView(vec_p, viewer);
2209:     VecView(vec_eta_element, viewer);
2210:     PetscViewerDestroy(&viewer);
2211:   }

2213:   /* Dump vertex- or edge-based fields to a second .vtr file */
2214:   {
2215:     PetscViewer viewer;

2217:     PetscViewerVTKOpen(PetscObjectComm((PetscObject)da_eta_down_left), "ex4_down_left.vtr", FILE_MODE_WRITE, &viewer);
2218:     VecView(vec_eta_down_left, viewer);
2219:     VecView(vec_rho_down_left, viewer);
2220:     PetscViewerDestroy(&viewer);
2221:   }
2222:   if (ctx->dim == 3) {
2223:     PetscViewer viewer;

2225:     PetscViewerVTKOpen(PetscObjectComm((PetscObject)da_eta_back_left), "ex4_back_left.vtr", FILE_MODE_WRITE, &viewer);
2226:     VecView(vec_eta_back_left, viewer);
2227:     VecView(vec_rho_back_left, viewer);
2228:     PetscViewerDestroy(&viewer);
2229:   }
2230:   if (ctx->dim == 3) {
2231:     PetscViewer viewer;

2233:     PetscViewerVTKOpen(PetscObjectComm((PetscObject)da_eta_back_down), "ex4_back_down.vtr", FILE_MODE_WRITE, &viewer);
2234:     VecView(vec_eta_back_down, viewer);
2235:     VecView(vec_rho_back_down, viewer);
2236:     PetscViewerDestroy(&viewer);
2237:   }

2239:   /* Destroy DMDAs and Vecs */
2240:   VecDestroy(&vec_vel_avg);
2241:   VecDestroy(&vec_p);
2242:   VecDestroy(&vec_eta_element);
2243:   VecDestroy(&vec_eta_down_left);
2244:   if (ctx->dim == 3) {
2245:     VecDestroy(&vec_eta_back_left);
2246:     VecDestroy(&vec_eta_back_down);
2247:   }
2248:   VecDestroy(&vec_rho_down_left);
2249:   if (ctx->dim == 3) {
2250:     VecDestroy(&vec_rho_back_left);
2251:     VecDestroy(&vec_rho_back_down);
2252:   }
2253:   DMDestroy(&da_vel_avg);
2254:   DMDestroy(&da_p);
2255:   DMDestroy(&da_eta_element);
2256:   DMDestroy(&da_eta_down_left);
2257:   if (ctx->dim == 3) {
2258:     DMDestroy(&da_eta_back_left);
2259:     DMDestroy(&da_eta_back_down);
2260:   }
2261:   DMDestroy(&da_rho_down_left);
2262:   if (ctx->dim == 3) {
2263:     DMDestroy(&da_rho_back_left);
2264:     DMDestroy(&da_rho_back_down);
2265:   }
2266:   VecDestroy(&vel_avg);
2267:   DMDestroy(&dm_vel_avg);
2268:   return 0;
2269: }

2271: /*TEST

2273:    test:
2274:       suffix: direct_umfpack
2275:       requires: suitesparse !complex
2276:       nsize: 1
2277:       args: -dim 2 -coefficients layers -nondimensional 0 -stag_grid_x 12 -stag_grid_y 7 -pc_type lu -pc_factor_mat_solver_type umfpack -ksp_converged_reason

2279:    test:
2280:       suffix: direct_mumps
2281:       requires: mumps !complex
2282:       nsize: 9
2283:       args: -dim 2 -coefficients layers -nondimensional 0 -stag_grid_x 13 -stag_grid_y 8 -pc_type lu -pc_factor_mat_solver_type mumps -ksp_converged_reason

2285:    test:
2286:       suffix: isovisc_nondim_abf_mg
2287:       nsize: 1
2288:       args: -dim 2 -coefficients layers -nondimensional 1 -pc_type fieldsplit -pc_fieldsplit_type schur -ksp_converged_reason -fieldsplit_element_ksp_type preonly  -pc_fieldsplit_detect_saddle_point false -fieldsplit_face_pc_type mg -fieldsplit_face_pc_mg_levels 3 -stag_grid_x 24 -stag_grid_y 24 -fieldsplit_face_pc_mg_galerkin -fieldsplit_face_ksp_converged_reason -ksp_type fgmres -fieldsplit_element_pc_type none -fieldsplit_face_mg_levels_ksp_max_it 6 -pc_fieldsplit_schur_fact_type upper -isoviscous

2290:    test:
2291:       suffix: isovisc_nondim_abf_mg_2
2292:       nsize: 1
2293:       args: -dim 2 -coefficients layers -nondimensional -isoviscous -eta1 1.0 -stag_grid_x 32 -stag_grid_y 32 -ksp_type fgmres -pc_type fieldsplit -pc_fieldsplit_type schur -pc_fieldsplit_schur_fact_type upper -build_auxiliary_operator -fieldsplit_element_ksp_type preonly -fieldsplit_element_pc_type jacobi -fieldsplit_face_pc_type mg -fieldsplit_face_pc_mg_levels 3 -fieldsplit_face_pc_mg_galerkin -fieldsplit_face_mg_levels_pc_type jacobi -fieldsplit_face_mg_levels_ksp_type chebyshev -ksp_converged_reason

2295:    test:
2296:       suffix: nondim_abf_mg
2297:       requires: suitesparse !complex
2298:       nsize: 4
2299:       args: -dim 2 -coefficients layers -pc_type fieldsplit -pc_fieldsplit_type schur -ksp_converged_reason -fieldsplit_element_ksp_type preonly  -pc_fieldsplit_detect_saddle_point false -fieldsplit_face_pc_type mg -fieldsplit_face_pc_mg_levels 3 -fieldsplit_face_pc_mg_galerkin -fieldsplit_face_mg_coarse_pc_type redundant -fieldsplit_face_mg_coarse_redundant_pc_type lu -fieldsplit_face_mg_coarse_redundant_pc_factor_mat_solver_type umfpack  -ksp_type fgmres -fieldsplit_element_pc_type none -fieldsplit_face_mg_levels_ksp_max_it 6 -pc_fieldsplit_schur_fact_type upper -nondimensional -eta1 1e-2 -eta2 1.0 -ksp_monitor -fieldsplit_face_mg_levels_pc_type jacobi -fieldsplit_face_mg_levels_ksp_type gmres -fieldsplit_element_pc_type jacobi -pc_fieldsplit_schur_precondition selfp   -stag_grid_x 32 -stag_grid_y 32 -fieldsplit_face_ksp_monitor

2301:    test:
2302:       suffix: nondim_abf_lu
2303:       requires: suitesparse !complex
2304:       nsize: 1
2305:       args: -dim 2 -coefficients layers -pc_type fieldsplit -pc_fieldsplit_type schur -ksp_converged_reason -fieldsplit_element_ksp_type preonly  -pc_fieldsplit_detect_saddle_point false -ksp_type fgmres -fieldsplit_element_pc_type none -pc_fieldsplit_schur_fact_type upper -nondimensional -eta1 1e-2 -eta2 1.0 -isoviscous 0 -ksp_monitor -fieldsplit_element_pc_type jacobi -build_auxiliary_operator -fieldsplit_face_pc_type lu -fieldsplit_face_pc_factor_mat_solver_type umfpack -stag_grid_x 32 -stag_grid_y 32

2307:    test:
2308:       suffix: 3d_nondim_isovisc_abf_mg
2309:       requires: !single
2310:       nsize: 1
2311:       args: -dim 3 -coefficients layers -isoviscous -nondimensional -build_auxiliary_operator -pc_type fieldsplit -pc_fieldsplit_type schur -ksp_converged_reason -fieldsplit_element_ksp_type preonly  -pc_fieldsplit_detect_saddle_point false -fieldsplit_face_pc_type mg -fieldsplit_face_pc_mg_levels 3 -s 16 -fieldsplit_face_pc_mg_galerkin -fieldsplit_face_ksp_converged_reason -ksp_type fgmres -fieldsplit_element_pc_type none -fieldsplit_face_mg_levels_ksp_max_it 6 -pc_fieldsplit_schur_fact_type upper

2313:    test:
2314:       TODO: unstable across systems
2315:       suffix: monolithic
2316:       nsize: 1
2317:       requires: double !complex
2318:       args: -dim {{2 3}separate output} -s 16 -custom_pc_mat -pc_type mg -pc_mg_levels 3 -pc_mg_galerkin -mg_levels_ksp_type gmres -mg_levels_ksp_norm_type unpreconditioned -mg_levels_ksp_max_it 10 -mg_levels_pc_type jacobi -ksp_converged_reason

2320:    test:
2321:       suffix: 3d_nondim_isovisc_sinker_abf_mg
2322:       requires: !complex !single
2323:       nsize: 1
2324:       args: -dim 3 -coefficients sinker -isoviscous -nondimensional -pc_type fieldsplit -pc_fieldsplit_type schur -ksp_converged_reason -fieldsplit_element_ksp_type preonly  -pc_fieldsplit_detect_saddle_point false -fieldsplit_face_pc_type mg -fieldsplit_face_pc_mg_levels 3 -s 16 -fieldsplit_face_pc_mg_galerkin -fieldsplit_face_ksp_converged_reason -ksp_type fgmres -fieldsplit_element_pc_type none -fieldsplit_face_mg_levels_ksp_max_it 6 -pc_fieldsplit_schur_fact_type upper

2326:    test:
2327:       TODO: unstable across systems
2328:       suffix: 3d_nondim_mono_mg_lamemstyle
2329:       nsize: 1
2330:       requires: suitesparse
2331:       args: -dim 3 -coefficients layers -nondimensional -s 16 -custom_pc_mat -pc_type mg -pc_mg_galerkin -pc_mg_levels 2 -mg_levels_ksp_type richardson -mg_levels_pc_type jacobi -mg_levels_ksp_richardson_scale 0.5 -mg_levels_ksp_max_it 20 -mg_coarse_pc_type lu -mg_coarse_pc_factor_mat_solver_type umfpack -ksp_converged_reason

2333:    test:
2334:       suffix: 3d_isovisc_blob_cuda
2335:       requires: cuda
2336:       nsize: 2
2337:       args: -dim 3 -coefficients blob -isoviscous -s 8 -build_auxiliary_operator -fieldsplit_element_ksp_type preonly -fieldsplit_element_pc_type jacobi -fieldsplit_face_ksp_type fgmres -fieldsplit_face_mg_levels_ksp_max_it 6 -fieldsplit_face_mg_levels_ksp_norm_type none -fieldsplit_face_mg_levels_ksp_type chebyshev -fieldsplit_face_mg_levels_pc_type jacobi -fieldsplit_face_pc_mg_galerkin -fieldsplit_face_pc_mg_levels 3 -fieldsplit_face_pc_type mg -pc_fieldsplit_schur_fact_type upper -pc_fieldsplit_schur_precondition user -pc_fieldsplit_type schur -pc_type fieldsplit -dm_mat_type aijcusparse -dm_vec_type cuda -log_view -fieldsplit_face_pc_mg_log
2338:       filter: awk "/MGInterp/ {print \$NF}"

2340: TEST*/