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));
100: /* Create Linear Solver */
101: PetscCall(KSPCreate(PetscObjectComm((PetscObject)dm), &ksp));
102: PetscCall(KSPSetOperators(ksp, A, A));
104: /* Set Up Preconditioner */
105: {
106: IS is[2];
107: DMStagStencil stencil_not_element[10], stencil_element[1];
109: const char *name[2] = {"not_element", "element"};
111: PetscCall(KSPGetPC(ksp, &pc));
112: PetscCall(PCSetType(pc, PCFIELDSPLIT));
114: // First split is everything except elements (intentionally not provided in canonical order)
115: for (PetscInt c = 0; c < 4; ++c) {
116: stencil_not_element[c].loc = DMSTAG_BACK_DOWN_LEFT;
117: stencil_not_element[c].c = c;
118: }
119: stencil_not_element[4].loc = DMSTAG_LEFT;
120: stencil_not_element[4].c = 0;
121: stencil_not_element[5].loc = DMSTAG_BACK;
122: stencil_not_element[5].c = 0;
123: stencil_not_element[6].loc = DMSTAG_DOWN;
124: stencil_not_element[6].c = 0;
125: stencil_not_element[7].loc = DMSTAG_BACK_DOWN;
126: stencil_not_element[7].c = 0;
127: stencil_not_element[8].loc = DMSTAG_BACK_LEFT;
128: stencil_not_element[8].c = 0;
129: stencil_not_element[9].loc = DMSTAG_DOWN_LEFT;
130: stencil_not_element[9].c = 0;
132: // Second split is elements
133: stencil_element[0].loc = DMSTAG_ELEMENT;
134: stencil_element[0].c = 0;
136: PetscCall(DMStagCreateISFromStencils(dm, 10, stencil_not_element, &is[0]));
137: PetscCall(DMStagCreateISFromStencils(dm, 1, stencil_element, &is[1]));
139: for (PetscInt i = 0; i < 2; ++i) PetscCall(PCFieldSplitSetIS(pc, name[i], is[i]));
141: for (PetscInt i = 0; i < 2; ++i) PetscCall(ISDestroy(&is[i]));
142: }
144: /* Logic below modifies the PC directly, so this is the last chance to change the solver
145: from the command line */
146: PetscCall(KSPSetFromOptions(ksp));
148: /* If the fieldsplit PC wasn't overridden, further split */
149: {
150: PCType pc_type;
151: PetscBool is_fieldsplit;
153: PetscCall(KSPGetPC(ksp, &pc));
154: PetscCall(PCGetType(pc, &pc_type));
155: PetscCall(PetscStrcmp(pc_type, PCFIELDSPLIT, &is_fieldsplit));
156: if (is_fieldsplit) {
157: PC pc_not_element, pc_not_vertex_first_three, pc_face_and_edge;
159: {
160: DM dm_not_element;
161: IS is[2];
162: KSP *sub_ksp;
163: PetscInt n_splits;
164: DMStagStencil stencil_vertex_first_three[3], stencil_not_vertex_first_three[7];
165: const char *name[2] = {"vertex_first_three", "not_vertex_first_three"};
167: PetscCall(PCSetUp(pc)); // Set up the Fieldsplit PC
168: PetscCall(PCFieldSplitGetSubKSP(pc, &n_splits, &sub_ksp));
169: PetscAssert(n_splits == 2, PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Expected a Fieldsplit PC with two fields");
170: PetscCall(KSPGetPC(sub_ksp[0], &pc_not_element)); // Select first sub-KSP
171: PetscCall(PCSetType(pc_not_element, PCFIELDSPLIT));
172: PetscCall(PetscFree(sub_ksp));
174: // A compatible DM for the second top-level split
175: PetscCall(DMStagCreateCompatibleDMStag(dm, 4, 1, 1, 0, &dm_not_element));
177: // First split within not_element is vertex_first_three
178: for (PetscInt c = 0; c < 3; ++c) {
179: stencil_vertex_first_three[c].loc = DMSTAG_BACK_DOWN_LEFT;
180: stencil_vertex_first_three[c].c = c;
181: }
183: // Second split within not_element is everything else
184: stencil_not_vertex_first_three[0].loc = DMSTAG_BACK_DOWN_LEFT;
185: stencil_not_vertex_first_three[0].c = 3;
186: stencil_not_vertex_first_three[1].loc = DMSTAG_LEFT;
187: stencil_not_vertex_first_three[1].c = 0;
188: stencil_not_vertex_first_three[2].loc = DMSTAG_BACK;
189: stencil_not_vertex_first_three[2].c = 0;
190: stencil_not_vertex_first_three[3].loc = DMSTAG_DOWN;
191: stencil_not_vertex_first_three[3].c = 0;
192: stencil_not_vertex_first_three[4].loc = DMSTAG_BACK_DOWN;
193: stencil_not_vertex_first_three[4].c = 0;
194: stencil_not_vertex_first_three[5].loc = DMSTAG_BACK_LEFT;
195: stencil_not_vertex_first_three[5].c = 0;
196: stencil_not_vertex_first_three[6].loc = DMSTAG_DOWN_LEFT;
197: stencil_not_vertex_first_three[6].c = 0;
199: PetscCall(DMStagCreateISFromStencils(dm_not_element, 3, stencil_vertex_first_three, &is[0]));
200: PetscCall(DMStagCreateISFromStencils(dm_not_element, 7, stencil_not_vertex_first_three, &is[1]));
202: for (PetscInt i = 0; i < 2; ++i) PetscCall(PCFieldSplitSetIS(pc_not_element, name[i], is[i]));
204: for (PetscInt i = 0; i < 2; ++i) PetscCall(ISDestroy(&is[i]));
205: PetscCall(DMDestroy(&dm_not_element));
206: }
208: // Further split the second split of the first split
209: {
210: DM dm_not_vertex_first_three;
211: PetscInt n_splits;
212: IS is[2];
213: KSP *sub_ksp;
214: DMStagStencil stencil_vertex_fourth[1], stencil_face_and_edge[6];
215: const char *name[2] = {"vertex_fourth", "face_and_edge"};
217: PetscCall(PCSetUp(pc_not_element)); // Set up the Fieldsplit PC
218: PetscCall(PCFieldSplitGetSubKSP(pc_not_element, &n_splits, &sub_ksp));
219: PetscAssert(n_splits == 2, PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Expected a Fieldsplit PC with two fields");
220: PetscCall(KSPGetPC(sub_ksp[1], &pc_not_vertex_first_three)); // Select second sub-KSP
221: PetscCall(PCSetType(pc_not_vertex_first_three, PCFIELDSPLIT));
222: PetscCall(PetscFree(sub_ksp));
224: PetscCall(DMStagCreateCompatibleDMStag(dm, 1, 1, 1, 0, &dm_not_vertex_first_three));
226: // First split is 4th vertex entry
227: stencil_vertex_fourth[0].loc = DMSTAG_BACK_DOWN_LEFT;
228: stencil_vertex_fourth[0].c = 3;
230: // Second split is faces and edges
231: stencil_face_and_edge[0].loc = DMSTAG_LEFT;
232: stencil_face_and_edge[0].c = 0;
233: stencil_face_and_edge[1].loc = DMSTAG_BACK;
234: stencil_face_and_edge[1].c = 0;
235: stencil_face_and_edge[2].loc = DMSTAG_DOWN;
236: stencil_face_and_edge[2].c = 0;
237: stencil_face_and_edge[3].loc = DMSTAG_BACK_DOWN;
238: stencil_face_and_edge[3].c = 0;
239: stencil_face_and_edge[4].loc = DMSTAG_BACK_LEFT;
240: stencil_face_and_edge[4].c = 0;
241: stencil_face_and_edge[5].loc = DMSTAG_DOWN_LEFT;
242: stencil_face_and_edge[5].c = 0;
244: PetscCall(DMStagCreateISFromStencils(dm_not_vertex_first_three, 1, stencil_vertex_fourth, &is[0]));
245: PetscCall(DMStagCreateISFromStencils(dm_not_vertex_first_three, 6, stencil_face_and_edge, &is[1]));
247: for (PetscInt i = 0; i < 2; ++i) PetscCall(PCFieldSplitSetIS(pc_not_vertex_first_three, name[i], is[i]));
249: for (PetscInt i = 0; i < 2; ++i) PetscCall(ISDestroy(&is[i]));
250: PetscCall(DMDestroy(&dm_not_vertex_first_three));
251: }
253: // Further split the second split of the second split of the first split
254: {
255: DM dm_face_and_edge;
256: PetscInt n_splits;
257: IS is[2];
258: KSP *sub_ksp;
259: DMStagStencil stencil_face[3], stencil_edge[3];
260: const char *name[2] = {"face", "edge"};
262: PetscCall(PCSetUp(pc_not_vertex_first_three)); // Set up the Fieldsplit PC
263: PetscCall(PCFieldSplitGetSubKSP(pc_not_vertex_first_three, &n_splits, &sub_ksp));
264: PetscAssert(n_splits == 2, PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Expected a Fieldsplit PC with two fields");
265: PetscCall(KSPGetPC(sub_ksp[1], &pc_face_and_edge)); // Select second sub-KSP
266: PetscCall(PCSetType(pc_face_and_edge, PCFIELDSPLIT));
267: PetscCall(PetscFree(sub_ksp));
269: PetscCall(DMStagCreateCompatibleDMStag(dm, 0, 1, 1, 0, &dm_face_and_edge));
271: // First split is faces
272: stencil_face[0].loc = DMSTAG_LEFT;
273: stencil_face[0].c = 0;
274: stencil_face[1].loc = DMSTAG_BACK;
275: stencil_face[1].c = 0;
276: stencil_face[2].loc = DMSTAG_DOWN;
277: stencil_face[2].c = 0;
279: // Second split is edges
280: stencil_edge[0].loc = DMSTAG_BACK_DOWN;
281: stencil_edge[0].c = 0;
282: stencil_edge[1].loc = DMSTAG_BACK_LEFT;
283: stencil_edge[1].c = 0;
284: stencil_edge[2].loc = DMSTAG_DOWN_LEFT;
285: stencil_edge[2].c = 0;
287: PetscCall(DMStagCreateISFromStencils(dm_face_and_edge, 3, stencil_face, &is[0]));
288: PetscCall(DMStagCreateISFromStencils(dm_face_and_edge, 3, stencil_edge, &is[1]));
290: for (PetscInt i = 0; i < 2; ++i) PetscCall(PCFieldSplitSetIS(pc_face_and_edge, name[i], is[i]));
292: for (PetscInt i = 0; i < 2; ++i) PetscCall(ISDestroy(&is[i]));
293: PetscCall(DMDestroy(&dm_face_and_edge));
294: }
295: }
296: }
298: /* Solve */
299: PetscCall(KSPSolve(ksp, b, x));
301: PetscCall(VecView(x, PETSC_VIEWER_STDOUT_WORLD));
303: /* Clean Up */
304: PetscCall(KSPDestroy(&ksp));
305: PetscCall(MatDestroy(&A));
306: PetscCall(VecDestroy(&x));
307: PetscCall(VecDestroy(&b));
308: PetscCall(DMDestroy(&dm));
309: PetscCall(PetscFinalize());
310: return 0;
311: }
313: /*TEST
315: test:
316: nsize: 8
317: 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
319: TEST*/