Actual source code: ex40.c

  1: static char help[] = "Test coloring for finite difference Jacobians with DMStag\n\n";

  3: #include <petscdm.h>
  4: #include <petscdmstag.h>
  5: #include <petscsnes.h>

  7: /*
  8:    Note that DMStagGetValuesStencil and DMStagSetValuesStencil are inefficient,
  9:    compared to DMStagVecGetArray() and friends, and only used here for testing
 10:    purposes, as they allow the code for the Jacobian and residual functions to
 11:    be more similar. In the intended application, where users are not writing
 12:    their own Jacobian assembly routines, one should use the faster, array-based
 13:    approach.
 14: */

 16: /* A "diagonal" objective function which only couples dof living at the same "point" */
 17: PetscErrorCode FormFunction1DNoCoupling(SNES snes, Vec x, Vec f, void *ctx)
 18: {
 19:   PetscInt start, n, n_extra, N, dof[2];
 20:   Vec      x_local;
 21:   DM       dm;

 23:   PetscFunctionBegin;
 24:   (void)ctx;
 25:   PetscCall(SNESGetDM(snes, &dm));
 26:   PetscCall(DMGetLocalVector(dm, &x_local));
 27:   PetscCall(DMGlobalToLocal(dm, x, INSERT_VALUES, x_local));
 28:   PetscCall(DMStagGetCorners(dm, &start, NULL, NULL, &n, NULL, NULL, &n_extra, NULL, NULL));
 29:   PetscCall(DMStagGetGlobalSizes(dm, &N, NULL, NULL));
 30:   PetscCall(DMStagGetDOF(dm, &dof[0], &dof[1], NULL, NULL));
 31:   for (PetscInt e = start; e < start + n + n_extra; ++e) {
 32:     for (PetscInt c = 0; c < dof[0]; ++c) {
 33:       DMStagStencil row;
 34:       PetscScalar   x_val, val;

 36:       row.i   = e;
 37:       row.loc = DMSTAG_LEFT;
 38:       row.c   = c;
 39:       PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
 40:       val = (10.0 + c) * x_val * x_val * x_val; // f_i = (10 +c) * x_i^3
 41:       PetscCall(DMStagVecSetValuesStencil(dm, f, 1, &row, &val, INSERT_VALUES));
 42:     }
 43:     if (e < N) {
 44:       for (PetscInt c = 0; c < dof[1]; ++c) {
 45:         DMStagStencil row;
 46:         PetscScalar   x_val, val;

 48:         row.i   = e;
 49:         row.loc = DMSTAG_ELEMENT;
 50:         row.c   = c;
 51:         PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
 52:         val = (20.0 + c) * x_val * x_val * x_val; // f_i = (20 + c) * x_i^3
 53:         PetscCall(DMStagVecSetValuesStencil(dm, f, 1, &row, &val, INSERT_VALUES));
 54:       }
 55:     }
 56:   }
 57:   PetscCall(DMRestoreLocalVector(dm, &x_local));
 58:   PetscCall(VecAssemblyBegin(f));
 59:   PetscCall(VecAssemblyEnd(f));
 60:   PetscFunctionReturn(PETSC_SUCCESS);
 61: }

 63: PetscErrorCode FormJacobian1DNoCoupling(SNES snes, Vec x, Mat Amat, Mat Pmat, void *ctx)
 64: {
 65:   PetscInt start, n, n_extra, N, dof[2];
 66:   Vec      x_local;
 67:   DM       dm;

 69:   PetscFunctionBegin;
 70:   (void)ctx;
 71:   PetscCall(SNESGetDM(snes, &dm));
 72:   PetscCall(DMGetLocalVector(dm, &x_local));
 73:   PetscCall(DMGlobalToLocal(dm, x, INSERT_VALUES, x_local));
 74:   PetscCall(DMStagGetCorners(dm, &start, NULL, NULL, &n, NULL, NULL, &n_extra, NULL, NULL));
 75:   PetscCall(DMStagGetGlobalSizes(dm, &N, NULL, NULL));
 76:   PetscCall(DMStagGetDOF(dm, &dof[0], &dof[1], NULL, NULL));
 77:   for (PetscInt e = start; e < start + n + n_extra; ++e) {
 78:     for (PetscInt c = 0; c < dof[0]; ++c) {
 79:       DMStagStencil row_vertex;
 80:       PetscScalar   x_val, val;

 82:       row_vertex.i   = e;
 83:       row_vertex.loc = DMSTAG_LEFT;
 84:       row_vertex.c   = c;
 85:       PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row_vertex, &x_val));
 86:       val = 3.0 * (10.0 + c) * x_val * x_val;
 87:       PetscCall(DMStagMatSetValuesStencil(dm, Amat, 1, &row_vertex, 1, &row_vertex, &val, INSERT_VALUES));
 88:     }
 89:     if (e < N) {
 90:       for (PetscInt c = 0; c < dof[1]; ++c) {
 91:         DMStagStencil row_element;
 92:         PetscScalar   x_val, val;

 94:         row_element.i   = e;
 95:         row_element.loc = DMSTAG_ELEMENT;
 96:         row_element.c   = c;
 97:         PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row_element, &x_val));
 98:         val = 3.0 * (20.0 + c) * x_val * x_val;
 99:         PetscCall(DMStagMatSetValuesStencil(dm, Amat, 1, &row_element, 1, &row_element, &val, INSERT_VALUES));
100:       }
101:     }
102:   }
103:   PetscCall(DMRestoreLocalVector(dm, &x_local));

105:   PetscCall(MatAssemblyBegin(Amat, MAT_FINAL_ASSEMBLY));
106:   PetscCall(MatAssemblyEnd(Amat, MAT_FINAL_ASSEMBLY));
107:   PetscCheck(Amat == Pmat, PETSC_COMM_WORLD, PETSC_ERR_SUP, "Not implemented for distinct Amat and Pmat");
108:   PetscFunctionReturn(PETSC_SUCCESS);
109: }

111: /* Objective functions which use the DM's stencil width. */
112: PetscErrorCode FormFunction1D(SNES snes, Vec x, Vec f, void *ctx)
113: {
114:   Vec               x_local;
115:   PetscInt          dim, stencil_width, start, n, n_extra, N, dof[2];
116:   DMStagStencilType stencil_type;
117:   DM                dm;

119:   PetscFunctionBegin;
120:   (void)ctx;
121:   PetscCall(SNESGetDM(snes, &dm));
122:   PetscCall(DMGetDimension(dm, &dim));
123:   PetscCheck(dim == 1, PETSC_COMM_WORLD, PETSC_ERR_SUP, "DM dimension must be 1");
124:   PetscCall(DMStagGetStencilType(dm, &stencil_type));
125:   PetscCheck(stencil_type == DMSTAG_STENCIL_STAR || stencil_type == DMSTAG_STENCIL_BOX, PETSC_COMM_WORLD, PETSC_ERR_SUP, "Only star and box stencils supported");
126:   PetscCall(DMStagGetStencilWidth(dm, &stencil_width));

128:   PetscCall(DMGetLocalVector(dm, &x_local));
129:   PetscCall(DMGlobalToLocal(dm, x, INSERT_VALUES, x_local));
130:   PetscCall(DMStagGetCorners(dm, &start, NULL, NULL, &n, NULL, NULL, &n_extra, NULL, NULL));
131:   PetscCall(DMStagGetGlobalSizes(dm, &N, NULL, NULL));
132:   PetscCall(DMStagGetDOF(dm, &dof[0], &dof[1], NULL, NULL));

134:   PetscCall(VecZeroEntries(f));

136:   for (PetscInt e = start; e < start + n + n_extra; ++e) {
137:     DMStagStencil row_vertex, row_element;

139:     row_vertex.i   = e;
140:     row_vertex.loc = DMSTAG_LEFT;

142:     row_element.i   = e;
143:     row_element.loc = DMSTAG_ELEMENT;

145:     for (PetscInt offset = -stencil_width; offset <= stencil_width; ++offset) {
146:       const PetscInt e_offset = e + offset;

148:       // vertex --> vertex
149:       if (e_offset >= 0 && e_offset < N + 1) { // Does not fully wrap in the periodic case
150:         DMStagStencil col;
151:         PetscScalar   x_val, val;

153:         for (PetscInt c_row = 0; c_row < dof[0]; ++c_row) {
154:           row_vertex.c = c_row;
155:           for (PetscInt c_col = 0; c_col < dof[0]; ++c_col) {
156:             col.c   = c_col;
157:             col.i   = e_offset;
158:             col.loc = DMSTAG_LEFT;
159:             PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &col, &x_val));
160:             val = (10.0 + offset) * x_val * x_val * x_val;
161:             PetscCall(DMStagVecSetValuesStencil(dm, f, 1, &row_vertex, &val, ADD_VALUES));
162:           }
163:         }
164:       }

166:       // element --> vertex
167:       if (e_offset >= 0 && e_offset < N) { // Does not fully wrap in the periodic case
168:         DMStagStencil col;
169:         PetscScalar   x_val, val;

171:         for (PetscInt c_row = 0; c_row < dof[0]; ++c_row) {
172:           row_vertex.c = c_row;
173:           for (PetscInt c_col = 0; c_col < dof[1]; ++c_col) {
174:             col.c   = c_col;
175:             col.i   = e_offset;
176:             col.loc = DMSTAG_ELEMENT;
177:             PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &col, &x_val));
178:             val = (15.0 + offset) * x_val * x_val * x_val;
179:             PetscCall(DMStagVecSetValuesStencil(dm, f, 1, &row_vertex, &val, ADD_VALUES));
180:           }
181:         }
182:       }

184:       if (e < N) {
185:         // vertex --> element
186:         if (e_offset >= 0 && e_offset < N + 1) { // Does not fully wrap in the periodic case
187:           DMStagStencil col;
188:           PetscScalar   x_val, val;

190:           for (PetscInt c_row = 0; c_row < dof[1]; ++c_row) {
191:             row_element.c = c_row;
192:             for (PetscInt c_col = 0; c_col < dof[0]; ++c_col) {
193:               col.c   = c_col;
194:               col.i   = e_offset;
195:               col.loc = DMSTAG_LEFT;
196:               PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &col, &x_val));
197:               val = (25.0 + offset) * x_val * x_val * x_val;
198:               PetscCall(DMStagVecSetValuesStencil(dm, f, 1, &row_element, &val, ADD_VALUES));
199:             }
200:           }
201:         }

203:         // element --> element
204:         if (e_offset >= 0 && e_offset < N) { // Does not fully wrap in the periodic case
205:           DMStagStencil col;
206:           PetscScalar   x_val, val;

208:           for (PetscInt c_row = 0; c_row < dof[1]; ++c_row) {
209:             row_element.c = c_row;
210:             for (PetscInt c_col = 0; c_col < dof[1]; ++c_col) {
211:               col.c   = c_col;
212:               col.i   = e_offset;
213:               col.loc = DMSTAG_ELEMENT;
214:               PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &col, &x_val));
215:               val = (20.0 + offset) * x_val * x_val * x_val;
216:               PetscCall(DMStagVecSetValuesStencil(dm, f, 1, &row_element, &val, ADD_VALUES));
217:             }
218:           }
219:         }
220:       }
221:     }
222:   }
223:   PetscCall(DMRestoreLocalVector(dm, &x_local));
224:   PetscCall(VecAssemblyBegin(f));
225:   PetscCall(VecAssemblyEnd(f));
226:   PetscFunctionReturn(PETSC_SUCCESS);
227: }

229: PetscErrorCode FormJacobian1D(SNES snes, Vec x, Mat Amat, Mat Pmat, void *ctx)
230: {
231:   Vec      x_local;
232:   PetscInt dim, stencil_width, start, n, n_extra, N, dof[2];
233:   DM       dm;

235:   PetscFunctionBegin;
236:   (void)ctx;
237:   PetscCall(SNESGetDM(snes, &dm));
238:   PetscCall(DMGetDimension(dm, &dim));
239:   PetscCheck(dim == 1, PETSC_COMM_WORLD, PETSC_ERR_SUP, "DM dimension must be 1");
240:   PetscCall(DMStagGetStencilWidth(dm, &stencil_width));

242:   PetscCall(DMGetLocalVector(dm, &x_local));
243:   PetscCall(DMGlobalToLocal(dm, x, INSERT_VALUES, x_local));
244:   PetscCall(DMStagGetCorners(dm, &start, NULL, NULL, &n, NULL, NULL, &n_extra, NULL, NULL));
245:   PetscCall(DMStagGetGlobalSizes(dm, &N, NULL, NULL));
246:   PetscCall(DMStagGetDOF(dm, &dof[0], &dof[1], NULL, NULL));

248:   PetscCall(MatZeroEntries(Amat));

250:   for (PetscInt e = start; e < start + n + n_extra; ++e) {
251:     DMStagStencil row_vertex, row_element;

253:     row_vertex.i   = e;
254:     row_vertex.loc = DMSTAG_LEFT;

256:     row_element.i   = e;
257:     row_element.loc = DMSTAG_ELEMENT;

259:     for (PetscInt offset = -stencil_width; offset <= stencil_width; ++offset) {
260:       const PetscInt e_offset = e + offset;

262:       // vertex --> vertex
263:       if (e_offset >= 0 && e_offset < N + 1) {
264:         DMStagStencil col;
265:         PetscScalar   x_val, val;

267:         for (PetscInt c_row = 0; c_row < dof[0]; ++c_row) {
268:           row_vertex.c = c_row;
269:           for (PetscInt c_col = 0; c_col < dof[0]; ++c_col) {
270:             col.c   = c_col;
271:             col.i   = e_offset;
272:             col.loc = DMSTAG_LEFT;
273:             PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &col, &x_val));
274:             val = 3.0 * (10.0 + offset) * x_val * x_val;
275:             PetscCall(DMStagMatSetValuesStencil(dm, Amat, 1, &row_vertex, 1, &col, &val, ADD_VALUES));
276:           }
277:         }
278:       }

280:       // element --> vertex
281:       if (e_offset >= 0 && e_offset < N) {
282:         DMStagStencil col;
283:         PetscScalar   x_val, val;

285:         for (PetscInt c_row = 0; c_row < dof[0]; ++c_row) {
286:           row_vertex.c = c_row;
287:           for (PetscInt c_col = 0; c_col < dof[1]; ++c_col) {
288:             col.c   = c_col;
289:             col.i   = e_offset;
290:             col.loc = DMSTAG_ELEMENT;
291:             PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &col, &x_val));
292:             val = 3.0 * (15.0 + offset) * x_val * x_val;
293:             PetscCall(DMStagMatSetValuesStencil(dm, Amat, 1, &row_vertex, 1, &col, &val, ADD_VALUES));
294:           }
295:         }
296:       }

298:       if (e < N) {
299:         // vertex --> element
300:         if (e_offset >= 0 && e_offset < N + 1) {
301:           DMStagStencil col;
302:           PetscScalar   x_val, val;

304:           for (PetscInt c_row = 0; c_row < dof[1]; ++c_row) {
305:             row_element.c = c_row;
306:             for (PetscInt c_col = 0; c_col < dof[0]; ++c_col) {
307:               col.c   = c_col;
308:               col.i   = e_offset;
309:               col.loc = DMSTAG_LEFT;
310:               PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &col, &x_val));
311:               val = 3.0 * (25.0 + offset) * x_val * x_val;
312:               PetscCall(DMStagMatSetValuesStencil(dm, Amat, 1, &row_element, 1, &col, &val, ADD_VALUES));
313:             }
314:           }
315:         }

317:         // element --> element
318:         if (e_offset >= 0 && e_offset < N) {
319:           DMStagStencil col;
320:           PetscScalar   x_val, val;

322:           for (PetscInt c_row = 0; c_row < dof[1]; ++c_row) {
323:             row_element.c = c_row;
324:             for (PetscInt c_col = 0; c_col < dof[1]; ++c_col) {
325:               col.c   = c_col;
326:               col.i   = e_offset;
327:               col.loc = DMSTAG_ELEMENT;
328:               PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &col, &x_val));
329:               val = 3.0 * (20.0 + offset) * x_val * x_val;
330:               PetscCall(DMStagMatSetValuesStencil(dm, Amat, 1, &row_element, 1, &col, &val, ADD_VALUES));
331:             }
332:           }
333:         }
334:       }
335:     }
336:   }
337:   PetscCall(DMRestoreLocalVector(dm, &x_local));
338:   PetscCall(MatAssemblyBegin(Amat, MAT_FINAL_ASSEMBLY));
339:   PetscCall(MatAssemblyEnd(Amat, MAT_FINAL_ASSEMBLY));
340:   PetscCheck(Amat == Pmat, PETSC_COMM_WORLD, PETSC_ERR_SUP, "Not implemented for distinct Amat and Pmat");
341:   PetscFunctionReturn(PETSC_SUCCESS);
342: }

344: PetscErrorCode FormFunction2DNoCoupling(SNES snes, Vec x, Vec f, void *ctx)
345: {
346:   PetscInt start[2], n[2], n_extra[2], N[2], dof[3];
347:   Vec      x_local;
348:   DM       dm;

350:   PetscFunctionBegin;
351:   (void)ctx;
352:   PetscCall(SNESGetDM(snes, &dm));
353:   PetscCall(DMGetLocalVector(dm, &x_local));
354:   PetscCall(DMGlobalToLocal(dm, x, INSERT_VALUES, x_local));
355:   PetscCall(DMStagGetCorners(dm, &start[0], &start[1], NULL, &n[0], &n[1], NULL, &n_extra[0], &n_extra[1], NULL));
356:   PetscCall(DMStagGetGlobalSizes(dm, &N[0], &N[1], NULL));
357:   PetscCall(DMStagGetDOF(dm, &dof[0], &dof[1], &dof[2], NULL));
358:   for (PetscInt ey = start[1]; ey < start[1] + n[1] + n_extra[1]; ++ey) {
359:     for (PetscInt ex = start[0]; ex < start[0] + n[0] + n_extra[0]; ++ex) {
360:       for (PetscInt c = 0; c < dof[0]; ++c) {
361:         DMStagStencil row;
362:         PetscScalar   x_val, val;

364:         row.i   = ex;
365:         row.j   = ey;
366:         row.loc = DMSTAG_DOWN_LEFT;
367:         row.c   = c;
368:         PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
369:         val = (5.0 + c) * x_val * x_val * x_val;
370:         PetscCall(DMStagVecSetValuesStencil(dm, f, 1, &row, &val, INSERT_VALUES));
371:       }
372:       if (ex < N[0]) {
373:         for (PetscInt c = 0; c < dof[1]; ++c) {
374:           DMStagStencil row;
375:           PetscScalar   x_val, val;

377:           row.i   = ex;
378:           row.j   = ey;
379:           row.loc = DMSTAG_DOWN;
380:           row.c   = c;
381:           PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
382:           val = (10.0 + c) * x_val * x_val * x_val;
383:           PetscCall(DMStagVecSetValuesStencil(dm, f, 1, &row, &val, INSERT_VALUES));
384:         }
385:       }
386:       if (ey < N[1]) {
387:         for (PetscInt c = 0; c < dof[1]; ++c) {
388:           DMStagStencil row;
389:           PetscScalar   x_val, val;

391:           row.i   = ex;
392:           row.j   = ey;
393:           row.loc = DMSTAG_LEFT;
394:           row.c   = c;
395:           PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
396:           val = (15.0 + c) * x_val * x_val * x_val;
397:           PetscCall(DMStagVecSetValuesStencil(dm, f, 1, &row, &val, INSERT_VALUES));
398:         }
399:       }
400:       if (ex < N[0] && ey < N[1]) {
401:         for (PetscInt c = 0; c < dof[2]; ++c) {
402:           DMStagStencil row;
403:           PetscScalar   x_val, val;

405:           row.i   = ex;
406:           row.j   = ey;
407:           row.loc = DMSTAG_ELEMENT;
408:           row.c   = c;
409:           PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
410:           val = (20.0 + c) * x_val * x_val * x_val;
411:           PetscCall(DMStagVecSetValuesStencil(dm, f, 1, &row, &val, INSERT_VALUES));
412:         }
413:       }
414:     }
415:   }
416:   PetscCall(DMRestoreLocalVector(dm, &x_local));
417:   PetscCall(VecAssemblyBegin(f));
418:   PetscCall(VecAssemblyEnd(f));
419:   PetscFunctionReturn(PETSC_SUCCESS);
420: }

422: PetscErrorCode FormJacobian2DNoCoupling(SNES snes, Vec x, Mat Amat, Mat Pmat, void *ctx)
423: {
424:   PetscInt start[2], n[2], n_extra[2], N[2], dof[3];
425:   Vec      x_local;
426:   DM       dm;

428:   PetscFunctionBegin;
429:   (void)ctx;
430:   PetscCall(SNESGetDM(snes, &dm));
431:   PetscCall(DMGetLocalVector(dm, &x_local));
432:   PetscCall(DMGlobalToLocal(dm, x, INSERT_VALUES, x_local));
433:   PetscCall(DMStagGetCorners(dm, &start[0], &start[1], NULL, &n[0], &n[1], NULL, &n_extra[0], &n_extra[1], NULL));
434:   PetscCall(DMStagGetGlobalSizes(dm, &N[0], &N[1], NULL));
435:   PetscCall(DMStagGetDOF(dm, &dof[0], &dof[1], &dof[2], NULL));
436:   for (PetscInt ey = start[1]; ey < start[1] + n[1] + n_extra[1]; ++ey) {
437:     for (PetscInt ex = start[0]; ex < start[0] + n[0] + n_extra[0]; ++ex) {
438:       for (PetscInt c = 0; c < dof[0]; ++c) {
439:         DMStagStencil row;
440:         PetscScalar   x_val, val;

442:         row.i   = ex;
443:         row.j   = ey;
444:         row.loc = DMSTAG_DOWN_LEFT;
445:         row.c   = c;
446:         PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
447:         val = 3.0 * (5.0 + c) * x_val * x_val;
448:         PetscCall(DMStagMatSetValuesStencil(dm, Amat, 1, &row, 1, &row, &val, INSERT_VALUES));
449:       }
450:       if (ex < N[0]) {
451:         for (PetscInt c = 0; c < dof[1]; ++c) {
452:           DMStagStencil row;
453:           PetscScalar   x_val, val;

455:           row.i   = ex;
456:           row.j   = ey;
457:           row.loc = DMSTAG_DOWN;
458:           row.c   = c;
459:           PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
460:           val = 3.0 * (10.0 + c) * x_val * x_val;
461:           PetscCall(DMStagMatSetValuesStencil(dm, Amat, 1, &row, 1, &row, &val, INSERT_VALUES));
462:         }
463:       }
464:       if (ey < N[1]) {
465:         for (PetscInt c = 0; c < dof[1]; ++c) {
466:           DMStagStencil row;
467:           PetscScalar   x_val, val;

469:           row.i   = ex;
470:           row.j   = ey;
471:           row.loc = DMSTAG_LEFT;
472:           row.c   = c;
473:           PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
474:           val = 3.0 * (15.0 + c) * x_val * x_val;
475:           PetscCall(DMStagMatSetValuesStencil(dm, Amat, 1, &row, 1, &row, &val, INSERT_VALUES));
476:         }
477:       }
478:       if (ex < N[0] && ey < N[1]) {
479:         for (PetscInt c = 0; c < dof[2]; ++c) {
480:           DMStagStencil row;
481:           PetscScalar   x_val, val;

483:           row.i   = ex;
484:           row.j   = ey;
485:           row.loc = DMSTAG_ELEMENT;
486:           row.c   = c;
487:           PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
488:           val = 3.0 * (20.0 + c) * x_val * x_val;
489:           PetscCall(DMStagMatSetValuesStencil(dm, Amat, 1, &row, 1, &row, &val, INSERT_VALUES));
490:         }
491:       }
492:     }
493:   }
494:   PetscCall(DMRestoreLocalVector(dm, &x_local));

496:   PetscCall(MatAssemblyBegin(Amat, MAT_FINAL_ASSEMBLY));
497:   PetscCall(MatAssemblyEnd(Amat, MAT_FINAL_ASSEMBLY));
498:   PetscCheck(Amat == Pmat, PETSC_COMM_WORLD, PETSC_ERR_SUP, "Not implemented for distinct Amat and Pmat");
499:   PetscFunctionReturn(PETSC_SUCCESS);
500: }

502: PetscErrorCode FormFunction2D(SNES snes, Vec x, Vec f, void *ctx)
503: {
504:   PetscInt start[2], n[2], n_extra[2], N[2], dof[3];
505:   Vec      x_local;
506:   DM       dm;

508:   PetscFunctionBegin;
509:   (void)ctx;
510:   PetscCall(SNESGetDM(snes, &dm));
511:   PetscCall(DMGetLocalVector(dm, &x_local));
512:   PetscCall(DMGlobalToLocal(dm, x, INSERT_VALUES, x_local));
513:   PetscCall(DMStagGetCorners(dm, &start[0], &start[1], NULL, &n[0], &n[1], NULL, &n_extra[0], &n_extra[1], NULL));
514:   PetscCall(DMStagGetGlobalSizes(dm, &N[0], &N[1], NULL));
515:   PetscCall(DMStagGetDOF(dm, &dof[0], &dof[1], &dof[2], NULL));

517:   PetscCall(VecZeroEntries(f));

519:   /* First, as in the simple case above */
520:   for (PetscInt ey = start[1]; ey < start[1] + n[1] + n_extra[1]; ++ey) {
521:     for (PetscInt ex = start[0]; ex < start[0] + n[0] + n_extra[0]; ++ex) {
522:       for (PetscInt c = 0; c < dof[0]; ++c) {
523:         DMStagStencil row;
524:         PetscScalar   x_val, val;

526:         row.i   = ex;
527:         row.j   = ey;
528:         row.loc = DMSTAG_DOWN_LEFT;
529:         row.c   = c;
530:         PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
531:         val = (5.0 + c) * x_val * x_val * x_val;
532:         PetscCall(DMStagVecSetValuesStencil(dm, f, 1, &row, &val, ADD_VALUES));
533:       }
534:       if (ex < N[0]) {
535:         for (PetscInt c = 0; c < dof[1]; ++c) {
536:           DMStagStencil row;
537:           PetscScalar   x_val, val;

539:           row.i   = ex;
540:           row.j   = ey;
541:           row.loc = DMSTAG_DOWN;
542:           row.c   = c;
543:           PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
544:           val = (10.0 + c) * x_val * x_val * x_val;
545:           PetscCall(DMStagVecSetValuesStencil(dm, f, 1, &row, &val, ADD_VALUES));
546:         }
547:       }
548:       if (ey < N[1]) {
549:         for (PetscInt c = 0; c < dof[1]; ++c) {
550:           DMStagStencil row;
551:           PetscScalar   x_val, val;

553:           row.i   = ex;
554:           row.j   = ey;
555:           row.loc = DMSTAG_LEFT;
556:           row.c   = c;
557:           PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
558:           val = (15.0 + c) * x_val * x_val * x_val;
559:           PetscCall(DMStagVecSetValuesStencil(dm, f, 1, &row, &val, ADD_VALUES));
560:         }
561:       }
562:       if (ex < N[0] && ey < N[1]) {
563:         for (PetscInt c = 0; c < dof[2]; ++c) {
564:           DMStagStencil row;
565:           PetscScalar   x_val, val;

567:           row.i   = ex;
568:           row.j   = ey;
569:           row.loc = DMSTAG_ELEMENT;
570:           row.c   = c;
571:           PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
572:           val = (20.0 + c) * x_val * x_val * x_val;
573:           PetscCall(DMStagVecSetValuesStencil(dm, f, 1, &row, &val, ADD_VALUES));
574:         }
575:       }
576:     }
577:   }

579:   /* Add additional terms fully coupling one interior element to another */
580:   {
581:     PetscMPIInt rank;

583:     PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)dm), &rank));
584:     if (rank == 0) {
585:       PetscInt       epe;
586:       DMStagStencil *row, *col;

588:       PetscCall(DMStagGetEntriesPerElement(dm, &epe));
589:       PetscCall(PetscMalloc1(epe, &row));
590:       PetscCall(PetscMalloc1(epe, &col));
591:       for (PetscInt i = 0; i < epe; ++i) {
592:         row[i].i = 0;
593:         row[i].j = 0;
594:         col[i].i = 0;
595:         col[i].j = 1;
596:       }
597:       {
598:         PetscInt nrows = 0;

600:         for (PetscInt c = 0; c < dof[0]; ++c) {
601:           row[nrows].c   = c;
602:           row[nrows].loc = DMSTAG_DOWN_LEFT;
603:           ++nrows;
604:         }
605:         for (PetscInt c = 0; c < dof[1]; ++c) {
606:           row[nrows].c   = c;
607:           row[nrows].loc = DMSTAG_LEFT;
608:           ++nrows;
609:         }
610:         for (PetscInt c = 0; c < dof[1]; ++c) {
611:           row[nrows].c   = c;
612:           row[nrows].loc = DMSTAG_DOWN;
613:           ++nrows;
614:         }
615:         for (PetscInt c = 0; c < dof[2]; ++c) {
616:           row[nrows].c   = c;
617:           row[nrows].loc = DMSTAG_ELEMENT;
618:           ++nrows;
619:         }
620:       }

622:       {
623:         PetscInt ncols = 0;

625:         for (PetscInt c = 0; c < dof[0]; ++c) {
626:           col[ncols].c   = c;
627:           col[ncols].loc = DMSTAG_DOWN_LEFT;
628:           ++ncols;
629:         }
630:         for (PetscInt c = 0; c < dof[1]; ++c) {
631:           col[ncols].c   = c;
632:           col[ncols].loc = DMSTAG_LEFT;
633:           ++ncols;
634:         }
635:         for (PetscInt c = 0; c < dof[1]; ++c) {
636:           col[ncols].c   = c;
637:           col[ncols].loc = DMSTAG_DOWN;
638:           ++ncols;
639:         }
640:         for (PetscInt c = 0; c < dof[2]; ++c) {
641:           col[ncols].c   = c;
642:           col[ncols].loc = DMSTAG_ELEMENT;
643:           ++ncols;
644:         }
645:       }

647:       for (PetscInt i = 0; i < epe; ++i) {
648:         for (PetscInt j = 0; j < epe; ++j) {
649:           PetscScalar x_val, val;

651:           PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &col[j], &x_val));
652:           val = (10 * i + j) * x_val * x_val * x_val;
653:           PetscCall(DMStagVecSetValuesStencil(dm, f, 1, &row[i], &val, ADD_VALUES));
654:         }
655:       }
656:       PetscCall(PetscFree(row));
657:       PetscCall(PetscFree(col));
658:     }
659:   }
660:   PetscCall(DMRestoreLocalVector(dm, &x_local));
661:   PetscCall(VecAssemblyBegin(f));
662:   PetscCall(VecAssemblyEnd(f));
663:   PetscFunctionReturn(PETSC_SUCCESS);
664: }

666: PetscErrorCode FormJacobian2D(SNES snes, Vec x, Mat Amat, Mat Pmat, void *ctx)
667: {
668:   PetscInt start[2], n[2], n_extra[2], N[2], dof[3];
669:   Vec      x_local;
670:   DM       dm;

672:   PetscFunctionBegin;
673:   (void)ctx;
674:   PetscCall(SNESGetDM(snes, &dm));
675:   PetscCall(DMGetLocalVector(dm, &x_local));
676:   PetscCall(DMGlobalToLocal(dm, x, INSERT_VALUES, x_local));
677:   PetscCall(DMStagGetCorners(dm, &start[0], &start[1], NULL, &n[0], &n[1], NULL, &n_extra[0], &n_extra[1], NULL));
678:   PetscCall(DMStagGetGlobalSizes(dm, &N[0], &N[1], NULL));
679:   PetscCall(DMStagGetDOF(dm, &dof[0], &dof[1], &dof[2], NULL));
680:   PetscCall(MatZeroEntries(Amat));
681:   for (PetscInt ey = start[1]; ey < start[1] + n[1] + n_extra[1]; ++ey) {
682:     for (PetscInt ex = start[0]; ex < start[0] + n[0] + n_extra[0]; ++ex) {
683:       for (PetscInt c = 0; c < dof[0]; ++c) {
684:         DMStagStencil row;
685:         PetscScalar   x_val, val;

687:         row.i   = ex;
688:         row.j   = ey;
689:         row.loc = DMSTAG_DOWN_LEFT;
690:         row.c   = c;
691:         PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
692:         val = 3.0 * (5.0 + c) * x_val * x_val;
693:         PetscCall(DMStagMatSetValuesStencil(dm, Amat, 1, &row, 1, &row, &val, ADD_VALUES));
694:       }
695:       if (ex < N[0]) {
696:         for (PetscInt c = 0; c < dof[1]; ++c) {
697:           DMStagStencil row;
698:           PetscScalar   x_val, val;

700:           row.i   = ex;
701:           row.j   = ey;
702:           row.loc = DMSTAG_DOWN;
703:           row.c   = c;
704:           PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
705:           val = 3.0 * (10.0 + c) * x_val * x_val;
706:           PetscCall(DMStagMatSetValuesStencil(dm, Amat, 1, &row, 1, &row, &val, ADD_VALUES));
707:         }
708:       }
709:       if (ey < N[1]) {
710:         for (PetscInt c = 0; c < dof[1]; ++c) {
711:           DMStagStencil row;
712:           PetscScalar   x_val, val;

714:           row.i   = ex;
715:           row.j   = ey;
716:           row.loc = DMSTAG_LEFT;
717:           row.c   = c;
718:           PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
719:           val = 3.0 * (15.0 + c) * x_val * x_val;
720:           PetscCall(DMStagMatSetValuesStencil(dm, Amat, 1, &row, 1, &row, &val, ADD_VALUES));
721:         }
722:       }
723:       if (ex < N[0] && ey < N[1]) {
724:         for (PetscInt c = 0; c < dof[2]; ++c) {
725:           DMStagStencil row;
726:           PetscScalar   x_val, val;

728:           row.i   = ex;
729:           row.j   = ey;
730:           row.loc = DMSTAG_ELEMENT;
731:           row.c   = c;
732:           PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
733:           val = 3.0 * (20.0 + c) * x_val * x_val;
734:           PetscCall(DMStagMatSetValuesStencil(dm, Amat, 1, &row, 1, &row, &val, ADD_VALUES));
735:         }
736:       }
737:     }
738:   }

740:   /* Add additional terms fully coupling one interior element to another */
741:   {
742:     PetscMPIInt rank;

744:     PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)dm), &rank));
745:     if (rank == 0) {
746:       PetscInt       epe;
747:       DMStagStencil *row, *col;

749:       PetscCall(DMStagGetEntriesPerElement(dm, &epe));
750:       PetscCall(PetscMalloc1(epe, &row));
751:       PetscCall(PetscMalloc1(epe, &col));
752:       for (PetscInt i = 0; i < epe; ++i) {
753:         row[i].i = 0;
754:         row[i].j = 0;
755:         col[i].i = 0;
756:         col[i].j = 1;
757:       }
758:       {
759:         PetscInt nrows = 0;

761:         for (PetscInt c = 0; c < dof[0]; ++c) {
762:           row[nrows].c   = c;
763:           row[nrows].loc = DMSTAG_DOWN_LEFT;
764:           ++nrows;
765:         }
766:         for (PetscInt c = 0; c < dof[1]; ++c) {
767:           row[nrows].c   = c;
768:           row[nrows].loc = DMSTAG_LEFT;
769:           ++nrows;
770:         }
771:         for (PetscInt c = 0; c < dof[1]; ++c) {
772:           row[nrows].c   = c;
773:           row[nrows].loc = DMSTAG_DOWN;
774:           ++nrows;
775:         }
776:         for (PetscInt c = 0; c < dof[2]; ++c) {
777:           row[nrows].c   = c;
778:           row[nrows].loc = DMSTAG_ELEMENT;
779:           ++nrows;
780:         }
781:       }

783:       {
784:         PetscInt ncols = 0;

786:         for (PetscInt c = 0; c < dof[0]; ++c) {
787:           col[ncols].c   = c;
788:           col[ncols].loc = DMSTAG_DOWN_LEFT;
789:           ++ncols;
790:         }
791:         for (PetscInt c = 0; c < dof[1]; ++c) {
792:           col[ncols].c   = c;
793:           col[ncols].loc = DMSTAG_LEFT;
794:           ++ncols;
795:         }
796:         for (PetscInt c = 0; c < dof[1]; ++c) {
797:           col[ncols].c   = c;
798:           col[ncols].loc = DMSTAG_DOWN;
799:           ++ncols;
800:         }
801:         for (PetscInt c = 0; c < dof[2]; ++c) {
802:           col[ncols].c   = c;
803:           col[ncols].loc = DMSTAG_ELEMENT;
804:           ++ncols;
805:         }
806:       }

808:       for (PetscInt i = 0; i < epe; ++i) {
809:         for (PetscInt j = 0; j < epe; ++j) {
810:           PetscScalar x_val, val;

812:           PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &col[j], &x_val));
813:           val = 3.0 * (10 * i + j) * x_val * x_val;
814:           PetscCall(DMStagMatSetValuesStencil(dm, Amat, 1, &row[i], 1, &col[j], &val, ADD_VALUES));
815:         }
816:       }
817:       PetscCall(PetscFree(row));
818:       PetscCall(PetscFree(col));
819:     }
820:   }
821:   PetscCall(DMRestoreLocalVector(dm, &x_local));
822:   PetscCall(MatAssemblyBegin(Amat, MAT_FINAL_ASSEMBLY));
823:   PetscCall(MatAssemblyEnd(Amat, MAT_FINAL_ASSEMBLY));
824:   PetscCheck(Amat == Pmat, PETSC_COMM_WORLD, PETSC_ERR_SUP, "Not implemented for distinct Amat and Pmat");
825:   PetscFunctionReturn(PETSC_SUCCESS);
826: }

828: PetscErrorCode FormFunction3DNoCoupling(SNES snes, Vec x, Vec f, void *ctx)
829: {
830:   PetscInt start[3], n[3], n_extra[3], N[3], dof[4];
831:   Vec      x_local;
832:   DM       dm;

834:   PetscFunctionBegin;
835:   (void)ctx;
836:   PetscCall(SNESGetDM(snes, &dm));
837:   PetscCall(DMGetLocalVector(dm, &x_local));
838:   PetscCall(DMGlobalToLocal(dm, x, INSERT_VALUES, x_local));
839:   PetscCall(DMStagGetCorners(dm, &start[0], &start[1], &start[2], &n[0], &n[1], &n[2], &n_extra[0], &n_extra[1], &n_extra[2]));
840:   PetscCall(DMStagGetGlobalSizes(dm, &N[0], &N[1], &N[2]));
841:   PetscCall(DMStagGetDOF(dm, &dof[0], &dof[1], &dof[2], &dof[3]));
842:   for (PetscInt ez = start[2]; ez < start[2] + n[2] + n_extra[2]; ++ez) {
843:     for (PetscInt ey = start[1]; ey < start[1] + n[1] + n_extra[1]; ++ey) {
844:       for (PetscInt ex = start[0]; ex < start[0] + n[0] + n_extra[0]; ++ex) {
845:         for (PetscInt c = 0; c < dof[0]; ++c) {
846:           DMStagStencil row;
847:           PetscScalar   x_val, val;

849:           row.i   = ex;
850:           row.j   = ey;
851:           row.k   = ez;
852:           row.loc = DMSTAG_BACK_DOWN_LEFT;
853:           row.c   = c;
854:           PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
855:           val = (5.0 + c) * x_val * x_val * x_val;
856:           PetscCall(DMStagVecSetValuesStencil(dm, f, 1, &row, &val, INSERT_VALUES));
857:         }
858:         if (ez < N[2]) {
859:           for (PetscInt c = 0; c < dof[1]; ++c) {
860:             DMStagStencil row;
861:             PetscScalar   x_val, val;

863:             row.i   = ex;
864:             row.j   = ey;
865:             row.k   = ez;
866:             row.loc = DMSTAG_DOWN_LEFT;
867:             row.c   = c;
868:             PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
869:             val = (50.0 + c) * x_val * x_val * x_val;
870:             PetscCall(DMStagVecSetValuesStencil(dm, f, 1, &row, &val, INSERT_VALUES));
871:           }
872:         }
873:         if (ey < N[1]) {
874:           for (PetscInt c = 0; c < dof[1]; ++c) {
875:             DMStagStencil row;
876:             PetscScalar   x_val, val;

878:             row.i   = ex;
879:             row.j   = ey;
880:             row.k   = ez;
881:             row.loc = DMSTAG_BACK_LEFT;
882:             row.c   = c;
883:             PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
884:             val = (55.0 + c) * x_val * x_val * x_val;
885:             PetscCall(DMStagVecSetValuesStencil(dm, f, 1, &row, &val, INSERT_VALUES));
886:           }
887:         }
888:         if (ex < N[0]) {
889:           for (PetscInt c = 0; c < dof[1]; ++c) {
890:             DMStagStencil row;
891:             PetscScalar   x_val, val;

893:             row.i   = ex;
894:             row.j   = ey;
895:             row.k   = ez;
896:             row.loc = DMSTAG_BACK_DOWN;
897:             row.c   = c;
898:             PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
899:             val = (60.0 + c) * x_val * x_val * x_val;
900:             PetscCall(DMStagVecSetValuesStencil(dm, f, 1, &row, &val, INSERT_VALUES));
901:           }
902:         }
903:         if (ex < N[0] && ez < N[2]) {
904:           for (PetscInt c = 0; c < dof[2]; ++c) {
905:             DMStagStencil row;
906:             PetscScalar   x_val, val;

908:             row.i   = ex;
909:             row.j   = ey;
910:             row.k   = ez;
911:             row.loc = DMSTAG_DOWN;
912:             row.c   = c;
913:             PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
914:             val = (10.0 + c) * x_val * x_val * x_val;
915:             PetscCall(DMStagVecSetValuesStencil(dm, f, 1, &row, &val, INSERT_VALUES));
916:           }
917:         }
918:         if (ey < N[1] && ez < N[2]) {
919:           for (PetscInt c = 0; c < dof[2]; ++c) {
920:             DMStagStencil row;
921:             PetscScalar   x_val, val;

923:             row.i   = ex;
924:             row.j   = ey;
925:             row.k   = ez;
926:             row.loc = DMSTAG_LEFT;
927:             row.c   = c;
928:             PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
929:             val = (15.0 + c) * x_val * x_val * x_val;
930:             PetscCall(DMStagVecSetValuesStencil(dm, f, 1, &row, &val, INSERT_VALUES));
931:           }
932:         }
933:         if (ex < N[0] && ey < N[1]) {
934:           for (PetscInt c = 0; c < dof[2]; ++c) {
935:             DMStagStencil row;
936:             PetscScalar   x_val, val;

938:             row.i   = ex;
939:             row.j   = ey;
940:             row.k   = ez;
941:             row.loc = DMSTAG_BACK;
942:             row.c   = c;
943:             PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
944:             val = (15.0 + c) * x_val * x_val * x_val;
945:             PetscCall(DMStagVecSetValuesStencil(dm, f, 1, &row, &val, INSERT_VALUES));
946:           }
947:         }
948:         if (ex < N[0] && ey < N[1] && ez < N[2]) {
949:           for (PetscInt c = 0; c < dof[3]; ++c) {
950:             DMStagStencil row;
951:             PetscScalar   x_val, val;

953:             row.i   = ex;
954:             row.j   = ey;
955:             row.k   = ez;
956:             row.loc = DMSTAG_ELEMENT;
957:             row.c   = c;
958:             PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
959:             val = (20.0 + c) * x_val * x_val * x_val;
960:             PetscCall(DMStagVecSetValuesStencil(dm, f, 1, &row, &val, INSERT_VALUES));
961:           }
962:         }
963:       }
964:     }
965:   }
966:   PetscCall(DMRestoreLocalVector(dm, &x_local));
967:   PetscCall(VecAssemblyBegin(f));
968:   PetscCall(VecAssemblyEnd(f));
969:   PetscFunctionReturn(PETSC_SUCCESS);
970: }

972: PetscErrorCode FormJacobian3DNoCoupling(SNES snes, Vec x, Mat Amat, Mat Pmat, void *ctx)
973: {
974:   PetscInt start[3], n[3], n_extra[3], N[3], dof[4];
975:   Vec      x_local;
976:   DM       dm;

978:   PetscFunctionBegin;
979:   (void)ctx;
980:   PetscCall(SNESGetDM(snes, &dm));
981:   PetscCall(DMGetLocalVector(dm, &x_local));
982:   PetscCall(DMGlobalToLocal(dm, x, INSERT_VALUES, x_local));
983:   PetscCall(DMStagGetCorners(dm, &start[0], &start[1], &start[2], &n[0], &n[1], &n[2], &n_extra[0], &n_extra[1], &n_extra[2]));
984:   PetscCall(DMStagGetGlobalSizes(dm, &N[0], &N[1], &N[2]));
985:   PetscCall(DMStagGetDOF(dm, &dof[0], &dof[1], &dof[2], &dof[3]));
986:   for (PetscInt ez = start[2]; ez < start[2] + n[2] + n_extra[2]; ++ez) {
987:     for (PetscInt ey = start[1]; ey < start[1] + n[1] + n_extra[1]; ++ey) {
988:       for (PetscInt ex = start[0]; ex < start[0] + n[0] + n_extra[0]; ++ex) {
989:         for (PetscInt c = 0; c < dof[0]; ++c) {
990:           DMStagStencil row;
991:           PetscScalar   x_val, val;

993:           row.i   = ex;
994:           row.j   = ey;
995:           row.k   = ez;
996:           row.loc = DMSTAG_BACK_DOWN_LEFT;
997:           row.c   = c;
998:           PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
999:           val = 3.0 * (5.0 + c) * x_val * x_val;
1000:           PetscCall(DMStagMatSetValuesStencil(dm, Amat, 1, &row, 1, &row, &val, INSERT_VALUES));
1001:         }
1002:         if (ez < N[2]) {
1003:           for (PetscInt c = 0; c < dof[1]; ++c) {
1004:             DMStagStencil row;
1005:             PetscScalar   x_val, val;

1007:             row.i   = ex;
1008:             row.j   = ey;
1009:             row.k   = ez;
1010:             row.loc = DMSTAG_DOWN_LEFT;
1011:             row.c   = c;
1012:             PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
1013:             val = 3.0 * (50.0 + c) * x_val * x_val;
1014:             PetscCall(DMStagMatSetValuesStencil(dm, Amat, 1, &row, 1, &row, &val, INSERT_VALUES));
1015:           }
1016:         }
1017:         if (ey < N[1]) {
1018:           for (PetscInt c = 0; c < dof[1]; ++c) {
1019:             DMStagStencil row;
1020:             PetscScalar   x_val, val;

1022:             row.i   = ex;
1023:             row.j   = ey;
1024:             row.k   = ez;
1025:             row.loc = DMSTAG_BACK_LEFT;
1026:             row.c   = c;
1027:             PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
1028:             val = 3.0 * (55.0 + c) * x_val * x_val;
1029:             PetscCall(DMStagMatSetValuesStencil(dm, Amat, 1, &row, 1, &row, &val, INSERT_VALUES));
1030:           }
1031:         }
1032:         if (ex < N[0]) {
1033:           for (PetscInt c = 0; c < dof[1]; ++c) {
1034:             DMStagStencil row;
1035:             PetscScalar   x_val, val;

1037:             row.i   = ex;
1038:             row.j   = ey;
1039:             row.k   = ez;
1040:             row.loc = DMSTAG_BACK_DOWN;
1041:             row.c   = c;
1042:             PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
1043:             val = 3.0 * (60.0 + c) * x_val * x_val;
1044:             PetscCall(DMStagMatSetValuesStencil(dm, Amat, 1, &row, 1, &row, &val, INSERT_VALUES));
1045:           }
1046:         }
1047:         if (ex < N[0] && ez < N[2]) {
1048:           for (PetscInt c = 0; c < dof[2]; ++c) {
1049:             DMStagStencil row;
1050:             PetscScalar   x_val, val;

1052:             row.i   = ex;
1053:             row.j   = ey;
1054:             row.k   = ez;
1055:             row.loc = DMSTAG_DOWN;
1056:             row.c   = c;
1057:             PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
1058:             val = 3.0 * (10.0 + c) * x_val * x_val;
1059:             PetscCall(DMStagMatSetValuesStencil(dm, Amat, 1, &row, 1, &row, &val, INSERT_VALUES));
1060:           }
1061:         }
1062:         if (ey < N[1] && ez < N[2]) {
1063:           for (PetscInt c = 0; c < dof[2]; ++c) {
1064:             DMStagStencil row;
1065:             PetscScalar   x_val, val;

1067:             row.i   = ex;
1068:             row.j   = ey;
1069:             row.k   = ez;
1070:             row.loc = DMSTAG_LEFT;
1071:             row.c   = c;
1072:             PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
1073:             val = 3.0 * (15.0 + c) * x_val * x_val;
1074:             PetscCall(DMStagMatSetValuesStencil(dm, Amat, 1, &row, 1, &row, &val, INSERT_VALUES));
1075:           }
1076:         }
1077:         if (ex < N[0] && ey < N[1]) {
1078:           for (PetscInt c = 0; c < dof[2]; ++c) {
1079:             DMStagStencil row;
1080:             PetscScalar   x_val, val;

1082:             row.i   = ex;
1083:             row.j   = ey;
1084:             row.k   = ez;
1085:             row.loc = DMSTAG_BACK;
1086:             row.c   = c;
1087:             PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
1088:             val = 3.0 * (15.0 + c) * x_val * x_val;
1089:             PetscCall(DMStagMatSetValuesStencil(dm, Amat, 1, &row, 1, &row, &val, INSERT_VALUES));
1090:           }
1091:         }
1092:         if (ex < N[0] && ey < N[1] && ez < N[2]) {
1093:           for (PetscInt c = 0; c < dof[3]; ++c) {
1094:             DMStagStencil row;
1095:             PetscScalar   x_val, val;

1097:             row.i   = ex;
1098:             row.j   = ey;
1099:             row.k   = ez;
1100:             row.loc = DMSTAG_ELEMENT;
1101:             row.c   = c;
1102:             PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
1103:             val = 3.0 * (20.0 + c) * x_val * x_val;
1104:             PetscCall(DMStagMatSetValuesStencil(dm, Amat, 1, &row, 1, &row, &val, INSERT_VALUES));
1105:           }
1106:         }
1107:       }
1108:     }
1109:   }
1110:   PetscCall(DMRestoreLocalVector(dm, &x_local));
1111:   PetscCall(MatAssemblyBegin(Amat, MAT_FINAL_ASSEMBLY));
1112:   PetscCall(MatAssemblyEnd(Amat, MAT_FINAL_ASSEMBLY));
1113:   PetscCheck(Amat == Pmat, PETSC_COMM_WORLD, PETSC_ERR_SUP, "Not implemented for distinct Amat and Pmat");
1114:   PetscFunctionReturn(PETSC_SUCCESS);
1115: }

1117: PetscErrorCode FormFunction3D(SNES snes, Vec x, Vec f, void *ctx)
1118: {
1119:   PetscInt start[3], n[3], n_extra[3], N[3], dof[4];
1120:   Vec      x_local;
1121:   DM       dm;

1123:   PetscFunctionBegin;
1124:   (void)ctx;
1125:   PetscCall(SNESGetDM(snes, &dm));
1126:   PetscCall(DMGetLocalVector(dm, &x_local));
1127:   PetscCall(DMGlobalToLocal(dm, x, INSERT_VALUES, x_local));
1128:   PetscCall(DMStagGetCorners(dm, &start[0], &start[1], &start[2], &n[0], &n[1], &n[2], &n_extra[0], &n_extra[1], &n_extra[2]));
1129:   PetscCall(DMStagGetGlobalSizes(dm, &N[0], &N[1], &N[2]));
1130:   PetscCall(DMStagGetDOF(dm, &dof[0], &dof[1], &dof[2], &dof[3]));
1131:   PetscCall(VecZeroEntries(f));
1132:   for (PetscInt ez = start[2]; ez < start[2] + n[2] + n_extra[2]; ++ez) {
1133:     for (PetscInt ey = start[1]; ey < start[1] + n[1] + n_extra[1]; ++ey) {
1134:       for (PetscInt ex = start[0]; ex < start[0] + n[0] + n_extra[0]; ++ex) {
1135:         for (PetscInt c = 0; c < dof[0]; ++c) {
1136:           DMStagStencil row;
1137:           PetscScalar   x_val, val;

1139:           row.i   = ex;
1140:           row.j   = ey;
1141:           row.k   = ez;
1142:           row.loc = DMSTAG_BACK_DOWN_LEFT;
1143:           row.c   = c;
1144:           PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
1145:           val = (5.0 + c) * x_val * x_val * x_val;
1146:           PetscCall(DMStagVecSetValuesStencil(dm, f, 1, &row, &val, ADD_VALUES));
1147:         }
1148:         if (ez < N[2]) {
1149:           for (PetscInt c = 0; c < dof[1]; ++c) {
1150:             DMStagStencil row;
1151:             PetscScalar   x_val, val;

1153:             row.i   = ex;
1154:             row.j   = ey;
1155:             row.k   = ez;
1156:             row.loc = DMSTAG_DOWN_LEFT;
1157:             row.c   = c;
1158:             PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
1159:             val = (50.0 + c) * x_val * x_val * x_val;
1160:             PetscCall(DMStagVecSetValuesStencil(dm, f, 1, &row, &val, ADD_VALUES));
1161:           }
1162:         }
1163:         if (ey < N[1]) {
1164:           for (PetscInt c = 0; c < dof[1]; ++c) {
1165:             DMStagStencil row;
1166:             PetscScalar   x_val, val;

1168:             row.i   = ex;
1169:             row.j   = ey;
1170:             row.k   = ez;
1171:             row.loc = DMSTAG_BACK_LEFT;
1172:             row.c   = c;
1173:             PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
1174:             val = (55.0 + c) * x_val * x_val * x_val;
1175:             PetscCall(DMStagVecSetValuesStencil(dm, f, 1, &row, &val, ADD_VALUES));
1176:           }
1177:         }
1178:         if (ex < N[0]) {
1179:           for (PetscInt c = 0; c < dof[1]; ++c) {
1180:             DMStagStencil row;
1181:             PetscScalar   x_val, val;

1183:             row.i   = ex;
1184:             row.j   = ey;
1185:             row.k   = ez;
1186:             row.loc = DMSTAG_BACK_DOWN;
1187:             row.c   = c;
1188:             PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
1189:             val = (60.0 + c) * x_val * x_val * x_val;
1190:             PetscCall(DMStagVecSetValuesStencil(dm, f, 1, &row, &val, ADD_VALUES));
1191:           }
1192:         }
1193:         if (ex < N[0] && ez < N[2]) {
1194:           for (PetscInt c = 0; c < dof[2]; ++c) {
1195:             DMStagStencil row;
1196:             PetscScalar   x_val, val;

1198:             row.i   = ex;
1199:             row.j   = ey;
1200:             row.k   = ez;
1201:             row.loc = DMSTAG_DOWN;
1202:             row.c   = c;
1203:             PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
1204:             val = (10.0 + c) * x_val * x_val * x_val;
1205:             PetscCall(DMStagVecSetValuesStencil(dm, f, 1, &row, &val, ADD_VALUES));
1206:           }
1207:         }
1208:         if (ey < N[1] && ez < N[2]) {
1209:           for (PetscInt c = 0; c < dof[2]; ++c) {
1210:             DMStagStencil row;
1211:             PetscScalar   x_val, val;

1213:             row.i   = ex;
1214:             row.j   = ey;
1215:             row.k   = ez;
1216:             row.loc = DMSTAG_LEFT;
1217:             row.c   = c;
1218:             PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
1219:             val = (15.0 + c) * x_val * x_val * x_val;
1220:             PetscCall(DMStagVecSetValuesStencil(dm, f, 1, &row, &val, ADD_VALUES));
1221:           }
1222:         }
1223:         if (ex < N[0] && ey < N[1]) {
1224:           for (PetscInt c = 0; c < dof[2]; ++c) {
1225:             DMStagStencil row;
1226:             PetscScalar   x_val, val;

1228:             row.i   = ex;
1229:             row.j   = ey;
1230:             row.k   = ez;
1231:             row.loc = DMSTAG_BACK;
1232:             row.c   = c;
1233:             PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
1234:             val = (15.0 + c) * x_val * x_val * x_val;
1235:             PetscCall(DMStagVecSetValuesStencil(dm, f, 1, &row, &val, ADD_VALUES));
1236:           }
1237:         }
1238:         if (ex < N[0] && ey < N[1] && ez < N[2]) {
1239:           for (PetscInt c = 0; c < dof[3]; ++c) {
1240:             DMStagStencil row;
1241:             PetscScalar   x_val, val;

1243:             row.i   = ex;
1244:             row.j   = ey;
1245:             row.k   = ez;
1246:             row.loc = DMSTAG_ELEMENT;
1247:             row.c   = c;
1248:             PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
1249:             val = (20.0 + c) * x_val * x_val * x_val;
1250:             PetscCall(DMStagVecSetValuesStencil(dm, f, 1, &row, &val, ADD_VALUES));
1251:           }
1252:         }
1253:       }
1254:     }
1255:   }

1257:   /* Add additional terms fully coupling one interior element to another */
1258:   {
1259:     PetscMPIInt rank;

1261:     PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)dm), &rank));
1262:     if (rank == 0) {
1263:       PetscInt       epe;
1264:       DMStagStencil *row, *col;

1266:       PetscCall(DMStagGetEntriesPerElement(dm, &epe));
1267:       PetscCall(PetscMalloc1(epe, &row));
1268:       PetscCall(PetscMalloc1(epe, &col));
1269:       for (PetscInt i = 0; i < epe; ++i) {
1270:         row[i].i = 0;
1271:         row[i].j = 0;
1272:         row[i].k = 0;
1273:         col[i].i = 0;
1274:         col[i].j = 0;
1275:         col[i].k = 1;
1276:       }

1278:       {
1279:         PetscInt nrows = 0;

1281:         for (PetscInt c = 0; c < dof[0]; ++c) {
1282:           row[nrows].c   = c;
1283:           row[nrows].loc = DMSTAG_BACK_DOWN_LEFT;
1284:           ++nrows;
1285:         }
1286:         for (PetscInt c = 0; c < dof[1]; ++c) {
1287:           row[nrows].c   = c;
1288:           row[nrows].loc = DMSTAG_DOWN_LEFT;
1289:           ++nrows;
1290:         }
1291:         for (PetscInt c = 0; c < dof[1]; ++c) {
1292:           row[nrows].c   = c;
1293:           row[nrows].loc = DMSTAG_BACK_LEFT;
1294:           ++nrows;
1295:         }
1296:         for (PetscInt c = 0; c < dof[1]; ++c) {
1297:           row[nrows].c   = c;
1298:           row[nrows].loc = DMSTAG_BACK_DOWN;
1299:           ++nrows;
1300:         }
1301:         for (PetscInt c = 0; c < dof[2]; ++c) {
1302:           row[nrows].c   = c;
1303:           row[nrows].loc = DMSTAG_LEFT;
1304:           ++nrows;
1305:         }
1306:         for (PetscInt c = 0; c < dof[2]; ++c) {
1307:           row[nrows].c   = c;
1308:           row[nrows].loc = DMSTAG_DOWN;
1309:           ++nrows;
1310:         }
1311:         for (PetscInt c = 0; c < dof[2]; ++c) {
1312:           row[nrows].c   = c;
1313:           row[nrows].loc = DMSTAG_BACK;
1314:           ++nrows;
1315:         }
1316:         for (PetscInt c = 0; c < dof[3]; ++c) {
1317:           row[nrows].c   = c;
1318:           row[nrows].loc = DMSTAG_ELEMENT;
1319:           ++nrows;
1320:         }
1321:       }

1323:       {
1324:         PetscInt ncols = 0;

1326:         for (PetscInt c = 0; c < dof[0]; ++c) {
1327:           col[ncols].c   = c;
1328:           col[ncols].loc = DMSTAG_BACK_DOWN_LEFT;
1329:           ++ncols;
1330:         }
1331:         for (PetscInt c = 0; c < dof[1]; ++c) {
1332:           col[ncols].c   = c;
1333:           col[ncols].loc = DMSTAG_DOWN_LEFT;
1334:           ++ncols;
1335:         }
1336:         for (PetscInt c = 0; c < dof[1]; ++c) {
1337:           col[ncols].c   = c;
1338:           col[ncols].loc = DMSTAG_BACK_LEFT;
1339:           ++ncols;
1340:         }
1341:         for (PetscInt c = 0; c < dof[1]; ++c) {
1342:           col[ncols].c   = c;
1343:           col[ncols].loc = DMSTAG_BACK_DOWN;
1344:           ++ncols;
1345:         }
1346:         for (PetscInt c = 0; c < dof[2]; ++c) {
1347:           col[ncols].c   = c;
1348:           col[ncols].loc = DMSTAG_LEFT;
1349:           ++ncols;
1350:         }
1351:         for (PetscInt c = 0; c < dof[2]; ++c) {
1352:           col[ncols].c   = c;
1353:           col[ncols].loc = DMSTAG_DOWN;
1354:           ++ncols;
1355:         }
1356:         for (PetscInt c = 0; c < dof[2]; ++c) {
1357:           col[ncols].c   = c;
1358:           col[ncols].loc = DMSTAG_BACK;
1359:           ++ncols;
1360:         }
1361:         for (PetscInt c = 0; c < dof[3]; ++c) {
1362:           col[ncols].c   = c;
1363:           col[ncols].loc = DMSTAG_ELEMENT;
1364:           ++ncols;
1365:         }
1366:       }

1368:       for (PetscInt i = 0; i < epe; ++i) {
1369:         for (PetscInt j = 0; j < epe; ++j) {
1370:           PetscScalar x_val, val;

1372:           PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &col[j], &x_val));
1373:           val = (10 * i + j) * x_val * x_val * x_val;
1374:           PetscCall(DMStagVecSetValuesStencil(dm, f, 1, &row[i], &val, ADD_VALUES));
1375:         }
1376:       }
1377:       PetscCall(PetscFree(row));
1378:       PetscCall(PetscFree(col));
1379:     }
1380:   }
1381:   PetscCall(DMRestoreLocalVector(dm, &x_local));
1382:   PetscCall(VecAssemblyBegin(f));
1383:   PetscCall(VecAssemblyEnd(f));
1384:   PetscFunctionReturn(PETSC_SUCCESS);
1385: }

1387: PetscErrorCode FormJacobian3D(SNES snes, Vec x, Mat Amat, Mat Pmat, void *ctx)
1388: {
1389:   PetscInt start[3], n[3], n_extra[3], N[3], dof[4];
1390:   Vec      x_local;
1391:   DM       dm;

1393:   PetscFunctionBegin;
1394:   (void)ctx;
1395:   PetscCall(SNESGetDM(snes, &dm));
1396:   PetscCall(DMGetLocalVector(dm, &x_local));
1397:   PetscCall(DMGlobalToLocal(dm, x, INSERT_VALUES, x_local));
1398:   PetscCall(DMStagGetCorners(dm, &start[0], &start[1], &start[2], &n[0], &n[1], &n[2], &n_extra[0], &n_extra[1], &n_extra[2]));
1399:   PetscCall(DMStagGetGlobalSizes(dm, &N[0], &N[1], &N[2]));
1400:   PetscCall(DMStagGetDOF(dm, &dof[0], &dof[1], &dof[2], &dof[3]));
1401:   PetscCall(MatZeroEntries(Amat));
1402:   for (PetscInt ez = start[2]; ez < start[2] + n[2] + n_extra[2]; ++ez) {
1403:     for (PetscInt ey = start[1]; ey < start[1] + n[1] + n_extra[1]; ++ey) {
1404:       for (PetscInt ex = start[0]; ex < start[0] + n[0] + n_extra[0]; ++ex) {
1405:         for (PetscInt c = 0; c < dof[0]; ++c) {
1406:           DMStagStencil row;
1407:           PetscScalar   x_val, val;

1409:           row.i   = ex;
1410:           row.j   = ey;
1411:           row.k   = ez;
1412:           row.loc = DMSTAG_BACK_DOWN_LEFT;
1413:           row.c   = c;
1414:           PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
1415:           val = 3.0 * (5.0 + c) * x_val * x_val;
1416:           PetscCall(DMStagMatSetValuesStencil(dm, Amat, 1, &row, 1, &row, &val, ADD_VALUES));
1417:         }
1418:         if (ez < N[2]) {
1419:           for (PetscInt c = 0; c < dof[1]; ++c) {
1420:             DMStagStencil row;
1421:             PetscScalar   x_val, val;

1423:             row.i   = ex;
1424:             row.j   = ey;
1425:             row.k   = ez;
1426:             row.loc = DMSTAG_DOWN_LEFT;
1427:             row.c   = c;
1428:             PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
1429:             val = 3.0 * (50.0 + c) * x_val * x_val;
1430:             PetscCall(DMStagMatSetValuesStencil(dm, Amat, 1, &row, 1, &row, &val, ADD_VALUES));
1431:           }
1432:         }
1433:         if (ey < N[1]) {
1434:           for (PetscInt c = 0; c < dof[1]; ++c) {
1435:             DMStagStencil row;
1436:             PetscScalar   x_val, val;

1438:             row.i   = ex;
1439:             row.j   = ey;
1440:             row.k   = ez;
1441:             row.loc = DMSTAG_BACK_LEFT;
1442:             row.c   = c;
1443:             PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
1444:             val = 3.0 * (55.0 + c) * x_val * x_val;
1445:             PetscCall(DMStagMatSetValuesStencil(dm, Amat, 1, &row, 1, &row, &val, ADD_VALUES));
1446:           }
1447:         }
1448:         if (ex < N[0]) {
1449:           for (PetscInt c = 0; c < dof[1]; ++c) {
1450:             DMStagStencil row;
1451:             PetscScalar   x_val, val;

1453:             row.i   = ex;
1454:             row.j   = ey;
1455:             row.k   = ez;
1456:             row.loc = DMSTAG_BACK_DOWN;
1457:             row.c   = c;
1458:             PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
1459:             val = 3.0 * (60.0 + c) * x_val * x_val;
1460:             PetscCall(DMStagMatSetValuesStencil(dm, Amat, 1, &row, 1, &row, &val, ADD_VALUES));
1461:           }
1462:         }
1463:         if (ex < N[0] && ez < N[2]) {
1464:           for (PetscInt c = 0; c < dof[2]; ++c) {
1465:             DMStagStencil row;
1466:             PetscScalar   x_val, val;

1468:             row.i   = ex;
1469:             row.j   = ey;
1470:             row.k   = ez;
1471:             row.loc = DMSTAG_DOWN;
1472:             row.c   = c;
1473:             PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
1474:             val = 3.0 * (10.0 + c) * x_val * x_val;
1475:             PetscCall(DMStagMatSetValuesStencil(dm, Amat, 1, &row, 1, &row, &val, ADD_VALUES));
1476:           }
1477:         }
1478:         if (ey < N[1] && ez < N[2]) {
1479:           for (PetscInt c = 0; c < dof[2]; ++c) {
1480:             DMStagStencil row;
1481:             PetscScalar   x_val, val;

1483:             row.i   = ex;
1484:             row.j   = ey;
1485:             row.k   = ez;
1486:             row.loc = DMSTAG_LEFT;
1487:             row.c   = c;
1488:             PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
1489:             val = 3.0 * (15.0 + c) * x_val * x_val;
1490:             PetscCall(DMStagMatSetValuesStencil(dm, Amat, 1, &row, 1, &row, &val, ADD_VALUES));
1491:           }
1492:         }
1493:         if (ex < N[0] && ey < N[1]) {
1494:           for (PetscInt c = 0; c < dof[2]; ++c) {
1495:             DMStagStencil row;
1496:             PetscScalar   x_val, val;

1498:             row.i   = ex;
1499:             row.j   = ey;
1500:             row.k   = ez;
1501:             row.loc = DMSTAG_BACK;
1502:             row.c   = c;
1503:             PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
1504:             val = 3.0 * (15.0 + c) * x_val * x_val;
1505:             PetscCall(DMStagMatSetValuesStencil(dm, Amat, 1, &row, 1, &row, &val, ADD_VALUES));
1506:           }
1507:         }
1508:         if (ex < N[0] && ey < N[1] && ez < N[2]) {
1509:           for (PetscInt c = 0; c < dof[3]; ++c) {
1510:             DMStagStencil row;
1511:             PetscScalar   x_val, val;

1513:             row.i   = ex;
1514:             row.j   = ey;
1515:             row.k   = ez;
1516:             row.loc = DMSTAG_ELEMENT;
1517:             row.c   = c;
1518:             PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &row, &x_val));
1519:             val = 3.0 * (20.0 + c) * x_val * x_val;
1520:             PetscCall(DMStagMatSetValuesStencil(dm, Amat, 1, &row, 1, &row, &val, ADD_VALUES));
1521:           }
1522:         }
1523:       }
1524:     }
1525:   }

1527:   /* Add additional terms fully coupling one interior element to another */
1528:   {
1529:     PetscMPIInt rank;

1531:     PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)dm), &rank));
1532:     if (rank == 0) {
1533:       PetscInt       epe;
1534:       DMStagStencil *row, *col;

1536:       PetscCall(DMStagGetEntriesPerElement(dm, &epe));
1537:       PetscCall(PetscMalloc1(epe, &row));
1538:       PetscCall(PetscMalloc1(epe, &col));
1539:       for (PetscInt i = 0; i < epe; ++i) {
1540:         row[i].i = 0;
1541:         row[i].j = 0;
1542:         row[i].k = 0;
1543:         col[i].i = 0;
1544:         col[i].j = 0;
1545:         col[i].k = 1;
1546:       }

1548:       {
1549:         PetscInt nrows = 0;

1551:         for (PetscInt c = 0; c < dof[0]; ++c) {
1552:           row[nrows].c   = c;
1553:           row[nrows].loc = DMSTAG_BACK_DOWN_LEFT;
1554:           ++nrows;
1555:         }
1556:         for (PetscInt c = 0; c < dof[1]; ++c) {
1557:           row[nrows].c   = c;
1558:           row[nrows].loc = DMSTAG_DOWN_LEFT;
1559:           ++nrows;
1560:         }
1561:         for (PetscInt c = 0; c < dof[1]; ++c) {
1562:           row[nrows].c   = c;
1563:           row[nrows].loc = DMSTAG_BACK_LEFT;
1564:           ++nrows;
1565:         }
1566:         for (PetscInt c = 0; c < dof[1]; ++c) {
1567:           row[nrows].c   = c;
1568:           row[nrows].loc = DMSTAG_BACK_DOWN;
1569:           ++nrows;
1570:         }
1571:         for (PetscInt c = 0; c < dof[2]; ++c) {
1572:           row[nrows].c   = c;
1573:           row[nrows].loc = DMSTAG_LEFT;
1574:           ++nrows;
1575:         }
1576:         for (PetscInt c = 0; c < dof[2]; ++c) {
1577:           row[nrows].c   = c;
1578:           row[nrows].loc = DMSTAG_DOWN;
1579:           ++nrows;
1580:         }
1581:         for (PetscInt c = 0; c < dof[2]; ++c) {
1582:           row[nrows].c   = c;
1583:           row[nrows].loc = DMSTAG_BACK;
1584:           ++nrows;
1585:         }
1586:         for (PetscInt c = 0; c < dof[3]; ++c) {
1587:           row[nrows].c   = c;
1588:           row[nrows].loc = DMSTAG_ELEMENT;
1589:           ++nrows;
1590:         }
1591:       }

1593:       {
1594:         PetscInt ncols = 0;

1596:         for (PetscInt c = 0; c < dof[0]; ++c) {
1597:           col[ncols].c   = c;
1598:           col[ncols].loc = DMSTAG_BACK_DOWN_LEFT;
1599:           ++ncols;
1600:         }
1601:         for (PetscInt c = 0; c < dof[1]; ++c) {
1602:           col[ncols].c   = c;
1603:           col[ncols].loc = DMSTAG_DOWN_LEFT;
1604:           ++ncols;
1605:         }
1606:         for (PetscInt c = 0; c < dof[1]; ++c) {
1607:           col[ncols].c   = c;
1608:           col[ncols].loc = DMSTAG_BACK_LEFT;
1609:           ++ncols;
1610:         }
1611:         for (PetscInt c = 0; c < dof[1]; ++c) {
1612:           col[ncols].c   = c;
1613:           col[ncols].loc = DMSTAG_BACK_DOWN;
1614:           ++ncols;
1615:         }
1616:         for (PetscInt c = 0; c < dof[2]; ++c) {
1617:           col[ncols].c   = c;
1618:           col[ncols].loc = DMSTAG_LEFT;
1619:           ++ncols;
1620:         }
1621:         for (PetscInt c = 0; c < dof[2]; ++c) {
1622:           col[ncols].c   = c;
1623:           col[ncols].loc = DMSTAG_DOWN;
1624:           ++ncols;
1625:         }
1626:         for (PetscInt c = 0; c < dof[2]; ++c) {
1627:           col[ncols].c   = c;
1628:           col[ncols].loc = DMSTAG_BACK;
1629:           ++ncols;
1630:         }
1631:         for (PetscInt c = 0; c < dof[3]; ++c) {
1632:           col[ncols].c   = c;
1633:           col[ncols].loc = DMSTAG_ELEMENT;
1634:           ++ncols;
1635:         }
1636:       }

1638:       for (PetscInt i = 0; i < epe; ++i) {
1639:         for (PetscInt j = 0; j < epe; ++j) {
1640:           PetscScalar x_val, val;

1642:           PetscCall(DMStagVecGetValuesStencil(dm, x_local, 1, &col[j], &x_val));
1643:           val = 3.0 * (10 * i + j) * x_val * x_val;
1644:           PetscCall(DMStagMatSetValuesStencil(dm, Amat, 1, &row[i], 1, &col[j], &val, ADD_VALUES));
1645:         }
1646:       }
1647:       PetscCall(PetscFree(row));
1648:       PetscCall(PetscFree(col));
1649:     }
1650:   }
1651:   PetscCall(DMRestoreLocalVector(dm, &x_local));
1652:   PetscCall(MatAssemblyBegin(Amat, MAT_FINAL_ASSEMBLY));
1653:   PetscCall(MatAssemblyEnd(Amat, MAT_FINAL_ASSEMBLY));
1654:   PetscCheck(Amat == Pmat, PETSC_COMM_WORLD, PETSC_ERR_SUP, "Not implemented for distinct Amat and Pmat");
1655:   PetscFunctionReturn(PETSC_SUCCESS);
1656: }

1658: int main(int argc, char **argv)
1659: {
1660:   DM        dm;
1661:   PetscInt  dim;
1662:   PetscBool no_coupling;
1663:   Vec       x, b;
1664:   SNES      snes;

1666:   PetscFunctionBeginUser;
1667:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
1668:   dim = 3;
1669:   PetscCall(PetscOptionsGetInt(NULL, NULL, "-dim", &dim, NULL));
1670:   no_coupling = PETSC_FALSE;
1671:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-no_coupling", &no_coupling, NULL));

1673:   switch (dim) {
1674:   case 1:
1675:     PetscCall(DMStagCreate1d(PETSC_COMM_WORLD, DM_BOUNDARY_NONE, 4, 1, 1, DMSTAG_STENCIL_BOX, 1, NULL, &dm));
1676:     break;
1677:   case 2:
1678:     PetscCall(DMStagCreate2d(PETSC_COMM_WORLD, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, 4, 3, PETSC_DECIDE, PETSC_DECIDE, 1, 1, 1, DMSTAG_STENCIL_BOX, 1, NULL, NULL, &dm));
1679:     break;
1680:   case 3:
1681:     PetscCall(DMStagCreate3d(PETSC_COMM_WORLD, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, 4, 3, 3, PETSC_DECIDE, PETSC_DECIDE, PETSC_DECIDE, 1, 1, 1, 1, DMSTAG_STENCIL_BOX, 1, NULL, NULL, NULL, &dm));
1682:     break;
1683:   default:
1684:     SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_SUP, "Unsupported dimension %" PetscInt_FMT, dim);
1685:   }
1686:   PetscCall(DMSetFromOptions(dm));
1687:   PetscCall(DMSetUp(dm));

1689:   PetscCall(SNESCreate(PETSC_COMM_WORLD, &snes));
1690:   PetscCall(SNESSetDM(snes, dm));
1691:   if (no_coupling) {
1692:     switch (dim) {
1693:     case 1:
1694:       PetscCall(SNESSetFunction(snes, NULL, FormFunction1DNoCoupling, NULL));
1695:       PetscCall(SNESSetJacobian(snes, NULL, NULL, FormJacobian1DNoCoupling, NULL));
1696:       break;
1697:     case 2:
1698:       PetscCall(SNESSetFunction(snes, NULL, FormFunction2DNoCoupling, NULL));
1699:       PetscCall(SNESSetJacobian(snes, NULL, NULL, FormJacobian2DNoCoupling, NULL));
1700:       break;
1701:     case 3:
1702:       PetscCall(SNESSetFunction(snes, NULL, FormFunction3DNoCoupling, NULL));
1703:       PetscCall(SNESSetJacobian(snes, NULL, NULL, FormJacobian3DNoCoupling, NULL));
1704:       break;
1705:     default:
1706:       SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_SUP, "Unsupported dimension %" PetscInt_FMT, dim);
1707:     }
1708:   } else {
1709:     switch (dim) {
1710:     case 1:
1711:       PetscCall(SNESSetFunction(snes, NULL, FormFunction1D, NULL));
1712:       PetscCall(SNESSetJacobian(snes, NULL, NULL, FormJacobian1D, NULL));
1713:       break;
1714:     case 2:
1715:       PetscCall(SNESSetFunction(snes, NULL, FormFunction2D, NULL));
1716:       PetscCall(SNESSetJacobian(snes, NULL, NULL, FormJacobian2D, NULL));
1717:       break;
1718:     case 3:
1719:       PetscCall(SNESSetFunction(snes, NULL, FormFunction3D, NULL));
1720:       PetscCall(SNESSetJacobian(snes, NULL, NULL, FormJacobian3D, NULL));
1721:       break;
1722:     default:
1723:       SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_SUP, "Unsupported dimension %" PetscInt_FMT, dim);
1724:     }
1725:   }
1726:   PetscCall(SNESSetFromOptions(snes));

1728:   PetscCall(DMCreateGlobalVector(dm, &x));
1729:   PetscCall(VecDuplicate(x, &b));
1730:   PetscCall(VecSet(x, 2.0)); // Initial guess
1731:   PetscCall(VecSet(b, 0.0)); // RHS
1732:   PetscCall(SNESSolve(snes, b, x));

1734:   PetscCall(SNESDestroy(&snes));
1735:   PetscCall(VecDestroy(&x));
1736:   PetscCall(VecDestroy(&b));
1737:   PetscCall(DMDestroy(&dm));
1738:   PetscCall(PetscFinalize());
1739:   return 0;
1740: }

1742: /*TEST

1744:    test:
1745:       suffix: 1d_no_coupling
1746:       nsize: {{1 2}separate output}
1747:       args: -dim 1 -no_coupling -stag_stencil_type none -pc_type jacobi -snes_converged_reason -snes_test_jacobian -stag_dof_0 {{1 2}separate output} -stag_dof_1 {{1 2}separate output} -snes_max_it 2
1748:    test:
1749:       suffix: 1d_test_jac
1750:       nsize: {{1 2}separate output}
1751:       args: -dim 1 -stag_stencil_width {{0 1}separate output} -pc_type jacobi -snes_converged_reason -snes_test_jacobian -snes_max_it 2
1752:    test:
1753:       suffix: 1d_fd_coloring
1754:       nsize: {{1 2}separate output}
1755:       args: -dim 1 -stag_stencil_width {{0 1 2}separate output} -pc_type jacobi -snes_converged_reason -snes_fd_color -snes_fd_color_use_mat -mat_coloring_type {{natural sl}} -snes_max_it 2
1756:    test:
1757:       suffix: 1d_periodic
1758:       nsize: {{1 2}separate output}
1759:       args: -dim 1 -stag_boundary_type_x periodic -stag_stencil_width {{1 2}separate output} -pc_type jacobi -snes_converged_reason -snes_test_jacobian -snes_max_it 2
1760:    test:
1761:       suffix: 1d_multidof
1762:       nsize: 2
1763:       args: -dim 1 -stag_stencil_width 2 -stag_dof_0 2 -stag_dof_1 3 -pc_type jacobi -snes_converged_reason -snes_test_jacobian -snes_max_it 2
1764:    test:
1765:       suffix: 2d_no_coupling
1766:       nsize: {{1 4}separate output}
1767:       args: -dim 2 -no_coupling -stag_stencil_type none -pc_type jacobi -snes_test_jacobian -stag_dof_0 {{1 2}separate output} -stag_dof_1 {{1 2}separate output} -stag_dof_2 {{1 2}separate output} -snes_max_it 2
1768:    test:
1769:       suffix: 3d_no_coupling
1770:       nsize: 2
1771:       args: -dim 3 -no_coupling -stag_stencil_type none -pc_type jacobi -snes_test_jacobian -stag_dof_0 2 -stag_dof_1 2 -stag_dof_2 2 -stag_dof_3 2 -snes_max_it 2
1772:    test:
1773:       suffix: 2d_fd_coloring
1774:       nsize: {{1 2}separate output}
1775:       args: -dim 2 -stag_stencil_width {{1 2}separate output} -pc_type jacobi -snes_converged_reason -snes_fd_color -snes_fd_color_use_mat -stag_stencil_type {{star box}separate output} -snes_max_it 2
1776:    test:
1777:       suffix: 3d_fd_coloring
1778:       nsize: {{1 2}separate output}
1779:       args: -dim 3 -stag_stencil_width {{1 2}separate output} -pc_type jacobi -snes_converged_reason -snes_fd_color -snes_fd_color_use_mat -stag_stencil_type {{star box}separate output} -snes_max_it 2
1780: TEST*/