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;

 18:   PetscFunctionBeginUser;
 19:   PetscCall(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:   PetscCall(DMSetFromOptions(dm));
 33:   PetscCall(DMSetUp(dm));

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

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

 50:     PetscCall(KSPGetConvergedReason(ksp, &reason));
 51:     PetscCheck(reason >= 0, PETSC_COMM_WORLD, PETSC_ERR_CONV_FAILED, "Linear solve failed");
 52:   }

 54:   /* Destroy PETSc objects and finalize */
 55:   PetscCall(MatDestroy(&A));
 56:   PetscCall(VecDestroy(&x));
 57:   PetscCall(VecDestroy(&b));
 58:   PetscCall(KSPDestroy(&ksp));
 59:   PetscCall(DMDestroy(&dm));
 60:   PetscCall(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;

 70:   PetscFunctionBeginUser;

 72:   PetscCall(DMCreateMatrix(dm, pA));
 73:   A = *pA;
 74:   PetscCall(DMCreateGlobalVector(dm, pb));
 75:   b = *pb;
 76:   PetscCall(DMStagGetCorners(dm, &start, NULL, NULL, &n, NULL, NULL, &n_extra, NULL, NULL));
 77:   PetscCall(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 boundary conditions (Dirichlet) */
 89:       PetscScalar val;

 91:       val = 1.0;
 92:       PetscCall(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:       PetscCall(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:       PetscCall(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:       PetscCall(DMStagVecSetValuesStencil(dm, b, 1, &row, &f, INSERT_VALUES));
133:     }
134:   }

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

141:   PetscFunctionReturn(PETSC_SUCCESS);
142: }

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

148:   PetscFunctionBeginUser;
149:   PetscCall(DMSetMatrixPreallocateOnly(dm, PETSC_TRUE));
150:   PetscCall(DMGetDimension(dm, &dim));
151:   switch (dim) {
152:   case 1:
153:     PetscCall(AssembleSystem1DVertexCentered(dm, pA, pb));
154:     break;
155:   default:
156:     SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_SUP, "Unsupported dimension: %" PetscInt_FMT, dim);
157:   }
158:   PetscFunctionReturn(PETSC_SUCCESS);
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*/