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*/