Actual source code: ex43.c
1: static char help[] = "Test using nested field splits with DMStag()\n\n";
3: #include <petscdm.h>
4: #include <petscdmstag.h>
5: #include <petscksp.h>
7: static PetscErrorCode AssembleSystem(DM dm, Mat A, Vec b)
8: {
9: PetscInt start[3], n[3], n_extra[3];
10: DMStagStencil row[11];
11: PetscScalar val[11];
13: PetscFunctionBeginUser;
14: PetscCall(DMStagGetCorners(dm, &start[0], &start[1], &start[2], &n[0], &n[1], &n[2], &n_extra[0], &n_extra[1], &n_extra[2]));
16: // Corner diagonal entries 10-14
17: for (PetscInt c = 0; c < 4; ++c) {
18: row[c].loc = DMSTAG_BACK_DOWN_LEFT;
19: row[c].c = c;
20: val[c] = 10.0 + c;
21: }
23: // Element entries 20
24: row[4].loc = DMSTAG_ELEMENT;
25: row[4].c = 0;
26: val[4] = 20.0;
28: // Face entries 30-32
29: row[5].loc = DMSTAG_BACK;
30: row[5].c = 0;
31: val[5] = 30.0;
33: row[6].loc = DMSTAG_LEFT;
34: row[6].c = 0;
35: val[6] = 32.0;
37: row[7].loc = DMSTAG_DOWN;
38: row[7].c = 0;
39: val[7] = 31.0;
41: // Edge entries 40-42
42: row[8].loc = DMSTAG_BACK_DOWN;
43: row[8].c = 0;
44: val[8] = 40.0;
46: row[9].loc = DMSTAG_BACK_LEFT;
47: row[9].c = 0;
48: val[9] = 41.0;
50: row[10].loc = DMSTAG_DOWN_LEFT;
51: row[10].c = 0;
52: val[10] = 42.0;
54: for (PetscInt k = start[2]; k < start[2] + n[2] + n_extra[2]; ++k) {
55: for (PetscInt j = start[1]; j < start[1] + n[1] + n_extra[1]; ++j) {
56: for (PetscInt i = start[0]; i < start[0] + n[0] + n_extra[0]; ++i) {
57: for (PetscInt e = 0; e < 11; ++e) {
58: row[e].i = i;
59: row[e].j = j;
60: row[e].k = k;
61: PetscCall(DMStagMatSetValuesStencil(dm, A, 1, &row[e], 1, &row[e], &val[e], INSERT_VALUES));
62: }
63: }
64: }
65: }
66: PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
67: PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
68: PetscCall(MatGetDiagonal(A, b)); // Get the diagonal, so x should be a constant 1.0
69: PetscFunctionReturn(PETSC_SUCCESS);
70: }
72: int main(int argc, char **argv)
73: {
74: DM dm;
75: KSP ksp;
76: PC pc;
77: Mat A;
78: Vec b, x;
80: PetscInt dof[4] = {4, 1, 1, 1};
81: PetscInt N[3] = {2, 3, 2};
83: PetscFunctionBeginUser;
84: PetscCall(PetscInitialize(&argc, &argv, NULL, help));
86: /* Create DM */
87: PetscCall(DMStagCreate3d(PETSC_COMM_WORLD, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, N[0], N[1], N[2], PETSC_DECIDE, PETSC_DECIDE, PETSC_DECIDE, dof[0], dof[1], dof[2], dof[3], DMSTAG_STENCIL_BOX, 1, NULL, NULL, NULL, &dm));
88: PetscCall(DMSetFromOptions(dm));
89: PetscCall(DMStagGetDOF(dm, &dof[0], &dof[1], &dof[2], &dof[3]));
90: PetscCall(DMStagGetGlobalSizes(dm, &N[0], &N[1], &N[2]));
91: PetscCall(DMSetUp(dm));
93: /* Create System */
94: PetscCall(DMSetMatrixPreallocateOnly(dm, PETSC_TRUE));
95: PetscCall(DMCreateMatrix(dm, &A));
96: PetscCall(DMCreateGlobalVector(dm, &b));
97: PetscCall(AssembleSystem(dm, A, b));
98: PetscCall(VecDuplicate(b, &x));
99: PetscCall(VecSet(x, 0.0));
101: /* Create Linear Solver */
102: PetscCall(KSPCreate(PetscObjectComm((PetscObject)dm), &ksp));
103: PetscCall(KSPSetOperators(ksp, A, A));
105: /* Set Up Preconditioner */
106: {
107: IS is[2];
108: DMStagStencil stencil_not_element[10], stencil_element[1];
110: const char *name[2] = {"not_element", "element"};
112: PetscCall(KSPGetPC(ksp, &pc));
113: PetscCall(PCSetType(pc, PCFIELDSPLIT));
115: // First split is everything except elements (intentionally not provided in canonical order)
116: for (PetscInt c = 0; c < 4; ++c) {
117: stencil_not_element[c].loc = DMSTAG_BACK_DOWN_LEFT;
118: stencil_not_element[c].c = c;
119: }
120: stencil_not_element[4].loc = DMSTAG_LEFT;
121: stencil_not_element[4].c = 0;
122: stencil_not_element[5].loc = DMSTAG_BACK;
123: stencil_not_element[5].c = 0;
124: stencil_not_element[6].loc = DMSTAG_DOWN;
125: stencil_not_element[6].c = 0;
126: stencil_not_element[7].loc = DMSTAG_BACK_DOWN;
127: stencil_not_element[7].c = 0;
128: stencil_not_element[8].loc = DMSTAG_BACK_LEFT;
129: stencil_not_element[8].c = 0;
130: stencil_not_element[9].loc = DMSTAG_DOWN_LEFT;
131: stencil_not_element[9].c = 0;
133: // Second split is elements
134: stencil_element[0].loc = DMSTAG_ELEMENT;
135: stencil_element[0].c = 0;
137: PetscCall(DMStagCreateISFromStencils(dm, 10, stencil_not_element, &is[0]));
138: PetscCall(DMStagCreateISFromStencils(dm, 1, stencil_element, &is[1]));
140: for (PetscInt i = 0; i < 2; ++i) PetscCall(PCFieldSplitSetIS(pc, name[i], is[i]));
142: for (PetscInt i = 0; i < 2; ++i) PetscCall(ISDestroy(&is[i]));
143: }
145: /* Logic below modifies the PC directly, so this is the last chance to change the solver
146: from the command line */
147: PetscCall(KSPSetFromOptions(ksp));
149: /* If the fieldsplit PC wasn't overridden, further split */
150: {
151: PCType pc_type;
152: PetscBool is_fieldsplit;
154: PetscCall(KSPGetPC(ksp, &pc));
155: PetscCall(PCGetType(pc, &pc_type));
156: PetscCall(PetscStrcmp(pc_type, PCFIELDSPLIT, &is_fieldsplit));
157: if (is_fieldsplit) {
158: PC pc_not_element, pc_not_vertex_first_three, pc_face_and_edge;
160: {
161: DM dm_not_element;
162: IS is[2];
163: KSP *sub_ksp;
164: PetscInt n_splits;
165: DMStagStencil stencil_vertex_first_three[3], stencil_not_vertex_first_three[7];
166: const char *name[2] = {"vertex_first_three", "not_vertex_first_three"};
168: PetscCall(PCSetUp(pc)); // Set up the Fieldsplit PC
169: PetscCall(PCFieldSplitGetSubKSP(pc, &n_splits, &sub_ksp));
170: PetscAssert(n_splits == 2, PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Expected a Fieldsplit PC with two fields");
171: PetscCall(KSPGetPC(sub_ksp[0], &pc_not_element)); // Select first sub-KSP
172: PetscCall(PCSetType(pc_not_element, PCFIELDSPLIT));
173: PetscCall(PetscFree(sub_ksp));
175: // A compatible DM for the second top-level split
176: PetscCall(DMStagCreateCompatibleDMStag(dm, 4, 1, 1, 0, &dm_not_element));
178: // First split within not_element is vertex_first_three
179: for (PetscInt c = 0; c < 3; ++c) {
180: stencil_vertex_first_three[c].loc = DMSTAG_BACK_DOWN_LEFT;
181: stencil_vertex_first_three[c].c = c;
182: }
184: // Second split within not_element is everything else
185: stencil_not_vertex_first_three[0].loc = DMSTAG_BACK_DOWN_LEFT;
186: stencil_not_vertex_first_three[0].c = 3;
187: stencil_not_vertex_first_three[1].loc = DMSTAG_LEFT;
188: stencil_not_vertex_first_three[1].c = 0;
189: stencil_not_vertex_first_three[2].loc = DMSTAG_BACK;
190: stencil_not_vertex_first_three[2].c = 0;
191: stencil_not_vertex_first_three[3].loc = DMSTAG_DOWN;
192: stencil_not_vertex_first_three[3].c = 0;
193: stencil_not_vertex_first_three[4].loc = DMSTAG_BACK_DOWN;
194: stencil_not_vertex_first_three[4].c = 0;
195: stencil_not_vertex_first_three[5].loc = DMSTAG_BACK_LEFT;
196: stencil_not_vertex_first_three[5].c = 0;
197: stencil_not_vertex_first_three[6].loc = DMSTAG_DOWN_LEFT;
198: stencil_not_vertex_first_three[6].c = 0;
200: PetscCall(DMStagCreateISFromStencils(dm_not_element, 3, stencil_vertex_first_three, &is[0]));
201: PetscCall(DMStagCreateISFromStencils(dm_not_element, 7, stencil_not_vertex_first_three, &is[1]));
203: for (PetscInt i = 0; i < 2; ++i) PetscCall(PCFieldSplitSetIS(pc_not_element, name[i], is[i]));
205: for (PetscInt i = 0; i < 2; ++i) PetscCall(ISDestroy(&is[i]));
206: PetscCall(DMDestroy(&dm_not_element));
207: }
209: // Further split the second split of the first split
210: {
211: DM dm_not_vertex_first_three;
212: PetscInt n_splits;
213: IS is[2];
214: KSP *sub_ksp;
215: DMStagStencil stencil_vertex_fourth[1], stencil_face_and_edge[6];
216: const char *name[2] = {"vertex_fourth", "face_and_edge"};
218: PetscCall(PCSetUp(pc_not_element)); // Set up the Fieldsplit PC
219: PetscCall(PCFieldSplitGetSubKSP(pc_not_element, &n_splits, &sub_ksp));
220: PetscAssert(n_splits == 2, PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Expected a Fieldsplit PC with two fields");
221: PetscCall(KSPGetPC(sub_ksp[1], &pc_not_vertex_first_three)); // Select second sub-KSP
222: PetscCall(PCSetType(pc_not_vertex_first_three, PCFIELDSPLIT));
223: PetscCall(PetscFree(sub_ksp));
225: PetscCall(DMStagCreateCompatibleDMStag(dm, 1, 1, 1, 0, &dm_not_vertex_first_three));
227: // First split is 4th vertex entry
228: stencil_vertex_fourth[0].loc = DMSTAG_BACK_DOWN_LEFT;
229: stencil_vertex_fourth[0].c = 3;
231: // Second split is faces and edges
232: stencil_face_and_edge[0].loc = DMSTAG_LEFT;
233: stencil_face_and_edge[0].c = 0;
234: stencil_face_and_edge[1].loc = DMSTAG_BACK;
235: stencil_face_and_edge[1].c = 0;
236: stencil_face_and_edge[2].loc = DMSTAG_DOWN;
237: stencil_face_and_edge[2].c = 0;
238: stencil_face_and_edge[3].loc = DMSTAG_BACK_DOWN;
239: stencil_face_and_edge[3].c = 0;
240: stencil_face_and_edge[4].loc = DMSTAG_BACK_LEFT;
241: stencil_face_and_edge[4].c = 0;
242: stencil_face_and_edge[5].loc = DMSTAG_DOWN_LEFT;
243: stencil_face_and_edge[5].c = 0;
245: PetscCall(DMStagCreateISFromStencils(dm_not_vertex_first_three, 1, stencil_vertex_fourth, &is[0]));
246: PetscCall(DMStagCreateISFromStencils(dm_not_vertex_first_three, 6, stencil_face_and_edge, &is[1]));
248: for (PetscInt i = 0; i < 2; ++i) PetscCall(PCFieldSplitSetIS(pc_not_vertex_first_three, name[i], is[i]));
250: for (PetscInt i = 0; i < 2; ++i) PetscCall(ISDestroy(&is[i]));
251: PetscCall(DMDestroy(&dm_not_vertex_first_three));
252: }
254: // Further split the second split of the second split of the first split
255: {
256: DM dm_face_and_edge;
257: PetscInt n_splits;
258: IS is[2];
259: KSP *sub_ksp;
260: DMStagStencil stencil_face[3], stencil_edge[3];
261: const char *name[2] = {"face", "edge"};
263: PetscCall(PCSetUp(pc_not_vertex_first_three)); // Set up the Fieldsplit PC
264: PetscCall(PCFieldSplitGetSubKSP(pc_not_vertex_first_three, &n_splits, &sub_ksp));
265: PetscAssert(n_splits == 2, PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Expected a Fieldsplit PC with two fields");
266: PetscCall(KSPGetPC(sub_ksp[1], &pc_face_and_edge)); // Select second sub-KSP
267: PetscCall(PCSetType(pc_face_and_edge, PCFIELDSPLIT));
268: PetscCall(PetscFree(sub_ksp));
270: PetscCall(DMStagCreateCompatibleDMStag(dm, 0, 1, 1, 0, &dm_face_and_edge));
272: // First split is faces
273: stencil_face[0].loc = DMSTAG_LEFT;
274: stencil_face[0].c = 0;
275: stencil_face[1].loc = DMSTAG_BACK;
276: stencil_face[1].c = 0;
277: stencil_face[2].loc = DMSTAG_DOWN;
278: stencil_face[2].c = 0;
280: // Second split is edges
281: stencil_edge[0].loc = DMSTAG_BACK_DOWN;
282: stencil_edge[0].c = 0;
283: stencil_edge[1].loc = DMSTAG_BACK_LEFT;
284: stencil_edge[1].c = 0;
285: stencil_edge[2].loc = DMSTAG_DOWN_LEFT;
286: stencil_edge[2].c = 0;
288: PetscCall(DMStagCreateISFromStencils(dm_face_and_edge, 3, stencil_face, &is[0]));
289: PetscCall(DMStagCreateISFromStencils(dm_face_and_edge, 3, stencil_edge, &is[1]));
291: for (PetscInt i = 0; i < 2; ++i) PetscCall(PCFieldSplitSetIS(pc_face_and_edge, name[i], is[i]));
293: for (PetscInt i = 0; i < 2; ++i) PetscCall(ISDestroy(&is[i]));
294: PetscCall(DMDestroy(&dm_face_and_edge));
295: }
296: }
297: }
299: /* Solve */
300: PetscCall(KSPSolve(ksp, b, x));
302: PetscCall(VecView(x, PETSC_VIEWER_STDOUT_WORLD));
304: /* Clean Up */
305: PetscCall(KSPDestroy(&ksp));
306: PetscCall(MatDestroy(&A));
307: PetscCall(VecDestroy(&x));
308: PetscCall(VecDestroy(&b));
309: PetscCall(DMDestroy(&dm));
310: PetscCall(PetscFinalize());
311: return 0;
312: }
314: /*TEST
316: test:
317: nsize: 8
318: args: -fieldsplit_element_ksp_max_it 1 -fieldsplit_element_ksp_type richardson -fieldsplit_element_pc_type none -fieldsplit_not_element_fieldsplit_not_vertex_first_three_fieldsplit_face_and_edge_fieldsplit_edge_ksp_max_it 1 -fieldsplit_not_element_fieldsplit_not_vertex_first_three_fieldsplit_face_and_edge_fieldsplit_edge_ksp_type richardson -fieldsplit_not_element_fieldsplit_not_vertex_first_three_fieldsplit_face_and_edge_fieldsplit_edge_pc_type none -fieldsplit_not_element_fieldsplit_not_vertex_first_three_fieldsplit_face_and_edge_fieldsplit_face_ksp_max_it 1 -fieldsplit_not_element_fieldsplit_not_vertex_first_three_fieldsplit_face_and_edge_fieldsplit_face_ksp_type richardson -fieldsplit_not_element_fieldsplit_not_vertex_first_three_fieldsplit_face_and_edge_fieldsplit_face_pc_type none -fieldsplit_not_element_fieldsplit_not_vertex_first_three_fieldsplit_face_and_edge_ksp_max_it 1 -fieldsplit_not_element_fieldsplit_not_vertex_first_three_fieldsplit_face_and_edge_ksp_type richardson -fieldsplit_not_element_fieldsplit_not_vertex_first_three_fieldsplit_face_and_edge_pc_fieldsplit_type additive -fieldsplit_not_element_fieldsplit_not_vertex_first_three_fieldsplit_face_and_edge_pc_type fieldsplit -fieldsplit_not_element_fieldsplit_not_vertex_first_three_fieldsplit_vertex_fourth_ksp_max_it 1 -fieldsplit_not_element_fieldsplit_not_vertex_first_three_fieldsplit_vertex_fourth_ksp_type richardson -fieldsplit_not_element_fieldsplit_not_vertex_first_three_fieldsplit_vertex_fourth_pc_type none -fieldsplit_not_element_fieldsplit_not_vertex_first_three_ksp_max_it 1 -fieldsplit_not_element_fieldsplit_not_vertex_first_three_ksp_type richardson -fieldsplit_not_element_fieldsplit_not_vertex_first_three_pc_fieldsplit_type additive -fieldsplit_not_element_fieldsplit_not_vertex_first_three_pc_type fieldsplit -fieldsplit_not_element_fieldsplit_vertex_first_three_ksp_max_it 1 -fieldsplit_not_element_fieldsplit_vertex_first_three_ksp_type richardson -fieldsplit_not_element_fieldsplit_vertex_first_three_pc_type none -fieldsplit_not_element_ksp_max_it 1 -fieldsplit_not_element_ksp_type richardson -fieldsplit_not_element_pc_fieldsplit_type additive -fieldsplit_not_element_pc_type fieldsplit -ksp_converged_reason -ksp_type preonly
320: TEST*/