1: static char help[] = "Solves the Poisson equation using DMStag, with a single field in 1D,\n"
  2:                      "intended to demonstrate the simplest use of geometric multigrid\n\n";

  4: #include <petscdm.h>
  5: #include <petscksp.h>
  6: #include <petscdmstag.h>

  8: static PetscErrorCode AssembleSystem(DM, Mat *, Vec *);

 10: int main(int argc, char **argv)
 11: {
 12:   Mat      A;
 13:   Vec      x, b;
 14:   KSP      ksp;
 15:   DM       dm;
 16:   PetscInt dim;

 19:   PetscInitialize(&argc, &argv, (char *)0, help);

 21:   dim = 1;

 23:   /* Create a DMStag object with a single degree of freedom for each point
 24:      on a single stratum, either vertices (0-cells) or elements (d-cells in d dimensions) */
 25:   if (dim == 1) {
 26:     PetscCall(DMStagCreate1d(PETSC_COMM_WORLD, DM_BOUNDARY_NONE, 8, /* Global element count */
 27:                              1,                                     /* Unknowns per vertex */
 28:                              0,                                     /* Unknowns per element */
 29:                              DMSTAG_STENCIL_BOX, 1,                 /* Elementwise stencil width */
 30:                              NULL, &dm));
 31:   } else SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_SUP, "Unsupported dimension: %" PetscInt_FMT, dim);
 32:   DMSetFromOptions(dm);
 33:   DMSetUp(dm);

 35:   /* Assemble the discrete system */
 36:   AssembleSystem(dm, &A, &b);

 38:   /* Solve */
 39:   KSPCreate(PetscObjectComm((PetscObject)dm), &ksp);
 40:   KSPSetType(ksp, KSPFGMRES);
 41:   KSPSetOperators(ksp, A, A);
 42:   KSPSetDM(ksp, dm);
 43:   KSPSetDMActive(ksp, PETSC_FALSE);
 44:   KSPSetFromOptions(ksp);
 45:   VecDuplicate(b, &x);
 46:   KSPSolve(ksp, b, x);
 47:   {
 48:     KSPConvergedReason reason;

 50:     KSPGetConvergedReason(ksp, &reason);
 52:   }

 54:   /* Destroy PETSc objects and finalize */
 55:   MatDestroy(&A);
 56:   VecDestroy(&x);
 57:   VecDestroy(&b);
 58:   KSPDestroy(&ksp);
 59:   DMDestroy(&dm);
 60:   PetscFinalize();
 61:   return 0;
 62: }

 64: static PetscErrorCode AssembleSystem1DVertexCentered(DM dm, Mat *pA, Vec *pb)
 65: {
 66:   Mat      A;
 67:   Vec      b;
 68:   PetscInt start, n, n_extra, N;


 72:   DMCreateMatrix(dm, pA);
 73:   A = *pA;
 74:   DMCreateGlobalVector(dm, pb);
 75:   b = *pb;
 76:   DMStagGetCorners(dm, &start, NULL, NULL, &n, NULL, NULL, &n_extra, NULL, NULL);
 77:   DMStagGetGlobalSizes(dm, &N, NULL, NULL);

 79:   /* Loop over all elements, including the non-physical, extra one past the right boundary */
 80:   for (PetscInt e = start; e < start + n + n_extra; ++e) {
 81:     DMStagStencil row;

 83:     row.i   = e;
 84:     row.c   = 0;
 85:     row.loc = DMSTAG_LEFT;

 87:     if (e == 0) {
 88:       /* Left bondary conditions (Dirichlet) */
 89:       PetscScalar val;

 91:       val = 1.0;
 92:       DMStagMatSetValuesStencil(dm, A, 1, &row, 1, &row, &val, INSERT_VALUES);
 93:     } else if (e == N) {
 94:       /* Right boundary (Dirichlet Boundary conditions)*/
 95:       DMStagStencil row_extra;
 96:       PetscScalar   val;

 98:       row_extra.i   = e;
 99:       row_extra.c   = 0;
100:       row_extra.loc = DMSTAG_LEFT;
101:       val           = 1.0;

103:       DMStagMatSetValuesStencil(dm, A, 1, &row_extra, 1, &row_extra, &val, INSERT_VALUES);
104:     } else {
105:       /* Interior */
106:       DMStagStencil col[3];
107:       PetscScalar   val[3];

109:       col[0].i   = e - 1;
110:       col[0].c   = 0;
111:       col[0].loc = DMSTAG_LEFT;
112:       val[0]     = 1.0;
113:       col[1].i   = e;
114:       col[1].c   = 0;
115:       col[1].loc = DMSTAG_LEFT;
116:       val[1]     = -2.0;
117:       col[2].i   = e + 1;
118:       col[2].c   = 0;
119:       col[2].loc = DMSTAG_LEFT;
120:       val[2]     = 1.0;

122:       DMStagMatSetValuesStencil(dm, A, 1, &row, 3, col, val, INSERT_VALUES);
123:     }

125:     /* Forcing */
126:     {
127:       PetscScalar x, f, h;

129:       h = 1.0 / N;              /* Assume a constant spacing instead of accessing coordinates */
130:       x = e / ((PetscScalar)N); // 0 - 1
131:       f = (x - 0.5) * h * h;    // Scale by h^2
132:       DMStagVecSetValuesStencil(dm, b, 1, &row, &f, INSERT_VALUES);
133:     }
134:   }

136:   MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY);
137:   MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY);
138:   VecAssemblyBegin(b);
139:   VecAssemblyEnd(b);

141:   return 0;
142: }

144: PetscErrorCode AssembleSystem(DM dm, Mat *pA, Vec *pb)
145: {
146:   PetscInt dim;

149:   DMSetMatrixPreallocateOnly(dm, PETSC_TRUE);
150:   DMGetDimension(dm, &dim);
151:   switch (dim) {
152:   case 1:
153:     AssembleSystem1DVertexCentered(dm, pA, pb);
154:     break;
155:   default:
156:     SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_SUP, "Unsupported dimension: %" PetscInt_FMT, dim);
157:   }
158:   return 0;
159: }

161: /*TEST
162:    test:
163:       args: -pc_type mg -pc_mg_galerkin -pc_mg_levels 7 -stag_grid_x 512 -ksp_converged_reason

165: TEST*/