1 static char help[] = "Test coloring for finite difference Jacobians with DMStag\n\n";
2
3 #include <petscdm.h>
4 #include <petscdmstag.h>
5 #include <petscsnes.h>
6
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 */
15
16 /* A "diagonal" objective function which only couples dof living at the same "point" */
FormFunction1DNoCoupling(SNES snes,Vec x,Vec f,PetscCtx ctx)17 PetscErrorCode FormFunction1DNoCoupling(SNES snes, Vec x, Vec f, PetscCtx ctx)
18 {
19 PetscInt start, n, n_extra, N, dof[2];
20 Vec x_local;
21 DM dm;
22
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;
35
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;
47
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 }
62
FormJacobian1DNoCoupling(SNES snes,Vec x,Mat Amat,Mat Pmat,PetscCtx ctx)63 PetscErrorCode FormJacobian1DNoCoupling(SNES snes, Vec x, Mat Amat, Mat Pmat, PetscCtx ctx)
64 {
65 PetscInt start, n, n_extra, N, dof[2];
66 Vec x_local;
67 DM dm;
68
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;
81
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;
93
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));
104
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 }
110
111 /* Objective functions which use the DM's stencil width. */
FormFunction1D(SNES snes,Vec x,Vec f,PetscCtx ctx)112 PetscErrorCode FormFunction1D(SNES snes, Vec x, Vec f, PetscCtx 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;
118
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));
127
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));
133
134 PetscCall(VecZeroEntries(f));
135
136 for (PetscInt e = start; e < start + n + n_extra; ++e) {
137 DMStagStencil row_vertex, row_element;
138
139 row_vertex.i = e;
140 row_vertex.loc = DMSTAG_LEFT;
141
142 row_element.i = e;
143 row_element.loc = DMSTAG_ELEMENT;
144
145 for (PetscInt offset = -stencil_width; offset <= stencil_width; ++offset) {
146 const PetscInt e_offset = e + offset;
147
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;
152
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 }
165
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;
170
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 }
183
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;
189
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 }
202
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;
207
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 }
228
FormJacobian1D(SNES snes,Vec x,Mat Amat,Mat Pmat,PetscCtx ctx)229 PetscErrorCode FormJacobian1D(SNES snes, Vec x, Mat Amat, Mat Pmat, PetscCtx ctx)
230 {
231 Vec x_local;
232 PetscInt dim, stencil_width, start, n, n_extra, N, dof[2];
233 DM dm;
234
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));
241
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));
247
248 PetscCall(MatZeroEntries(Amat));
249
250 for (PetscInt e = start; e < start + n + n_extra; ++e) {
251 DMStagStencil row_vertex, row_element;
252
253 row_vertex.i = e;
254 row_vertex.loc = DMSTAG_LEFT;
255
256 row_element.i = e;
257 row_element.loc = DMSTAG_ELEMENT;
258
259 for (PetscInt offset = -stencil_width; offset <= stencil_width; ++offset) {
260 const PetscInt e_offset = e + offset;
261
262 // vertex --> vertex
263 if (e_offset >= 0 && e_offset < N + 1) {
264 DMStagStencil col;
265 PetscScalar x_val, val;
266
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 }
279
280 // element --> vertex
281 if (e_offset >= 0 && e_offset < N) {
282 DMStagStencil col;
283 PetscScalar x_val, val;
284
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 }
297
298 if (e < N) {
299 // vertex --> element
300 if (e_offset >= 0 && e_offset < N + 1) {
301 DMStagStencil col;
302 PetscScalar x_val, val;
303
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 }
316
317 // element --> element
318 if (e_offset >= 0 && e_offset < N) {
319 DMStagStencil col;
320 PetscScalar x_val, val;
321
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 }
343
FormFunction2DNoCoupling(SNES snes,Vec x,Vec f,PetscCtx ctx)344 PetscErrorCode FormFunction2DNoCoupling(SNES snes, Vec x, Vec f, PetscCtx ctx)
345 {
346 PetscInt start[2], n[2], n_extra[2], N[2], dof[3];
347 Vec x_local;
348 DM dm;
349
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;
363
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;
376
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;
390
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;
404
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 }
421
FormJacobian2DNoCoupling(SNES snes,Vec x,Mat Amat,Mat Pmat,PetscCtx ctx)422 PetscErrorCode FormJacobian2DNoCoupling(SNES snes, Vec x, Mat Amat, Mat Pmat, PetscCtx ctx)
423 {
424 PetscInt start[2], n[2], n_extra[2], N[2], dof[3];
425 Vec x_local;
426 DM dm;
427
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;
441
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;
454
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;
468
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;
482
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));
495
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 }
501
FormFunction2D(SNES snes,Vec x,Vec f,PetscCtx ctx)502 PetscErrorCode FormFunction2D(SNES snes, Vec x, Vec f, PetscCtx ctx)
503 {
504 PetscInt start[2], n[2], n_extra[2], N[2], dof[3];
505 Vec x_local;
506 DM dm;
507
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));
516
517 PetscCall(VecZeroEntries(f));
518
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;
525
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;
538
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;
552
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;
566
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 }
578
579 /* Add additional terms fully coupling one interior element to another */
580 {
581 PetscMPIInt rank;
582
583 PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)dm), &rank));
584 if (rank == 0) {
585 PetscInt epe;
586 DMStagStencil *row, *col;
587
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;
599
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 }
621
622 {
623 PetscInt ncols = 0;
624
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 }
646
647 for (PetscInt i = 0; i < epe; ++i) {
648 for (PetscInt j = 0; j < epe; ++j) {
649 PetscScalar x_val, val;
650
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 }
665
FormJacobian2D(SNES snes,Vec x,Mat Amat,Mat Pmat,PetscCtx ctx)666 PetscErrorCode FormJacobian2D(SNES snes, Vec x, Mat Amat, Mat Pmat, PetscCtx ctx)
667 {
668 PetscInt start[2], n[2], n_extra[2], N[2], dof[3];
669 Vec x_local;
670 DM dm;
671
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;
686
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;
699
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;
713
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;
727
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 }
739
740 /* Add additional terms fully coupling one interior element to another */
741 {
742 PetscMPIInt rank;
743
744 PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)dm), &rank));
745 if (rank == 0) {
746 PetscInt epe;
747 DMStagStencil *row, *col;
748
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;
760
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 }
782
783 {
784 PetscInt ncols = 0;
785
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 }
807
808 for (PetscInt i = 0; i < epe; ++i) {
809 for (PetscInt j = 0; j < epe; ++j) {
810 PetscScalar x_val, val;
811
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 }
827
FormFunction3DNoCoupling(SNES snes,Vec x,Vec f,PetscCtx ctx)828 PetscErrorCode FormFunction3DNoCoupling(SNES snes, Vec x, Vec f, PetscCtx ctx)
829 {
830 PetscInt start[3], n[3], n_extra[3], N[3], dof[4];
831 Vec x_local;
832 DM dm;
833
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;
848
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;
862
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;
877
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;
892
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;
907
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;
922
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;
937
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;
952
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 }
971
FormJacobian3DNoCoupling(SNES snes,Vec x,Mat Amat,Mat Pmat,PetscCtx ctx)972 PetscErrorCode FormJacobian3DNoCoupling(SNES snes, Vec x, Mat Amat, Mat Pmat, PetscCtx ctx)
973 {
974 PetscInt start[3], n[3], n_extra[3], N[3], dof[4];
975 Vec x_local;
976 DM dm;
977
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;
992
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;
1006
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;
1021
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;
1036
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;
1051
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;
1066
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;
1081
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;
1096
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 }
1116
FormFunction3D(SNES snes,Vec x,Vec f,PetscCtx ctx)1117 PetscErrorCode FormFunction3D(SNES snes, Vec x, Vec f, PetscCtx ctx)
1118 {
1119 PetscInt start[3], n[3], n_extra[3], N[3], dof[4];
1120 Vec x_local;
1121 DM dm;
1122
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;
1138
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;
1152
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;
1167
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;
1182
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;
1197
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;
1212
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;
1227
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;
1242
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 }
1256
1257 /* Add additional terms fully coupling one interior element to another */
1258 {
1259 PetscMPIInt rank;
1260
1261 PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)dm), &rank));
1262 if (rank == 0) {
1263 PetscInt epe;
1264 DMStagStencil *row, *col;
1265
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 }
1277
1278 {
1279 PetscInt nrows = 0;
1280
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 }
1322
1323 {
1324 PetscInt ncols = 0;
1325
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 }
1367
1368 for (PetscInt i = 0; i < epe; ++i) {
1369 for (PetscInt j = 0; j < epe; ++j) {
1370 PetscScalar x_val, val;
1371
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 }
1386
FormJacobian3D(SNES snes,Vec x,Mat Amat,Mat Pmat,PetscCtx ctx)1387 PetscErrorCode FormJacobian3D(SNES snes, Vec x, Mat Amat, Mat Pmat, PetscCtx ctx)
1388 {
1389 PetscInt start[3], n[3], n_extra[3], N[3], dof[4];
1390 Vec x_local;
1391 DM dm;
1392
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;
1408
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;
1422
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;
1437
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;
1452
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;
1467
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;
1482
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;
1497
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;
1512
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 }
1526
1527 /* Add additional terms fully coupling one interior element to another */
1528 {
1529 PetscMPIInt rank;
1530
1531 PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)dm), &rank));
1532 if (rank == 0) {
1533 PetscInt epe;
1534 DMStagStencil *row, *col;
1535
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 }
1547
1548 {
1549 PetscInt nrows = 0;
1550
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 }
1592
1593 {
1594 PetscInt ncols = 0;
1595
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 }
1637
1638 for (PetscInt i = 0; i < epe; ++i) {
1639 for (PetscInt j = 0; j < epe; ++j) {
1640 PetscScalar x_val, val;
1641
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 }
1657
main(int argc,char ** argv)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;
1665
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));
1672
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));
1688
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));
1727
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));
1733
1734 PetscCall(SNESDestroy(&snes));
1735 PetscCall(VecDestroy(&x));
1736 PetscCall(VecDestroy(&b));
1737 PetscCall(DMDestroy(&dm));
1738 PetscCall(PetscFinalize());
1739 return 0;
1740 }
1741
1742 /*TEST
1743
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*/
1781