1 static char help[] = "Test using nested field splits with DMStag()\n\n";
2
3 #include <petscdm.h>
4 #include <petscdmstag.h>
5 #include <petscksp.h>
6
AssembleSystem(DM dm,Mat A,Vec b)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];
12
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]));
15
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 }
22
23 // Element entries 20
24 row[4].loc = DMSTAG_ELEMENT;
25 row[4].c = 0;
26 val[4] = 20.0;
27
28 // Face entries 30-32
29 row[5].loc = DMSTAG_BACK;
30 row[5].c = 0;
31 val[5] = 30.0;
32
33 row[6].loc = DMSTAG_LEFT;
34 row[6].c = 0;
35 val[6] = 32.0;
36
37 row[7].loc = DMSTAG_DOWN;
38 row[7].c = 0;
39 val[7] = 31.0;
40
41 // Edge entries 40-42
42 row[8].loc = DMSTAG_BACK_DOWN;
43 row[8].c = 0;
44 val[8] = 40.0;
45
46 row[9].loc = DMSTAG_BACK_LEFT;
47 row[9].c = 0;
48 val[9] = 41.0;
49
50 row[10].loc = DMSTAG_DOWN_LEFT;
51 row[10].c = 0;
52 val[10] = 42.0;
53
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 }
71
main(int argc,char ** argv)72 int main(int argc, char **argv)
73 {
74 DM dm;
75 KSP ksp;
76 PC pc;
77 Mat A;
78 Vec b, x;
79
80 PetscInt dof[4] = {4, 1, 1, 1};
81 PetscInt N[3] = {2, 3, 2};
82
83 PetscFunctionBeginUser;
84 PetscCall(PetscInitialize(&argc, &argv, NULL, help));
85
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));
92
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));
100
101 /* Create Linear Solver */
102 PetscCall(KSPCreate(PetscObjectComm((PetscObject)dm), &ksp));
103 PetscCall(KSPSetOperators(ksp, A, A));
104
105 /* Set Up Preconditioner */
106 {
107 IS is[2];
108 DMStagStencil stencil_not_element[10], stencil_element[1];
109
110 const char *name[2] = {"not_element", "element"};
111
112 PetscCall(KSPGetPC(ksp, &pc));
113 PetscCall(PCSetType(pc, PCFIELDSPLIT));
114
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;
132
133 // Second split is elements
134 stencil_element[0].loc = DMSTAG_ELEMENT;
135 stencil_element[0].c = 0;
136
137 PetscCall(DMStagCreateISFromStencils(dm, 10, stencil_not_element, &is[0]));
138 PetscCall(DMStagCreateISFromStencils(dm, 1, stencil_element, &is[1]));
139
140 for (PetscInt i = 0; i < 2; ++i) PetscCall(PCFieldSplitSetIS(pc, name[i], is[i]));
141
142 for (PetscInt i = 0; i < 2; ++i) PetscCall(ISDestroy(&is[i]));
143 }
144
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));
148
149 /* If the fieldsplit PC wasn't overridden, further split */
150 {
151 PCType pc_type;
152 PetscBool is_fieldsplit;
153
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;
159
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"};
167
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));
174
175 // A compatible DM for the second top-level split
176 PetscCall(DMStagCreateCompatibleDMStag(dm, 4, 1, 1, 0, &dm_not_element));
177
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 }
183
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;
199
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]));
202
203 for (PetscInt i = 0; i < 2; ++i) PetscCall(PCFieldSplitSetIS(pc_not_element, name[i], is[i]));
204
205 for (PetscInt i = 0; i < 2; ++i) PetscCall(ISDestroy(&is[i]));
206 PetscCall(DMDestroy(&dm_not_element));
207 }
208
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"};
217
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));
224
225 PetscCall(DMStagCreateCompatibleDMStag(dm, 1, 1, 1, 0, &dm_not_vertex_first_three));
226
227 // First split is 4th vertex entry
228 stencil_vertex_fourth[0].loc = DMSTAG_BACK_DOWN_LEFT;
229 stencil_vertex_fourth[0].c = 3;
230
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;
244
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]));
247
248 for (PetscInt i = 0; i < 2; ++i) PetscCall(PCFieldSplitSetIS(pc_not_vertex_first_three, name[i], is[i]));
249
250 for (PetscInt i = 0; i < 2; ++i) PetscCall(ISDestroy(&is[i]));
251 PetscCall(DMDestroy(&dm_not_vertex_first_three));
252 }
253
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"};
262
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));
269
270 PetscCall(DMStagCreateCompatibleDMStag(dm, 0, 1, 1, 0, &dm_face_and_edge));
271
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;
279
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;
287
288 PetscCall(DMStagCreateISFromStencils(dm_face_and_edge, 3, stencil_face, &is[0]));
289 PetscCall(DMStagCreateISFromStencils(dm_face_and_edge, 3, stencil_edge, &is[1]));
290
291 for (PetscInt i = 0; i < 2; ++i) PetscCall(PCFieldSplitSetIS(pc_face_and_edge, name[i], is[i]));
292
293 for (PetscInt i = 0; i < 2; ++i) PetscCall(ISDestroy(&is[i]));
294 PetscCall(DMDestroy(&dm_face_and_edge));
295 }
296 }
297 }
298
299 /* Solve */
300 PetscCall(KSPSolve(ksp, b, x));
301
302 PetscCall(VecView(x, PETSC_VIEWER_STDOUT_WORLD));
303
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 }
313
314 /*TEST
315
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
319
320 TEST*/
321