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, NULL, 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;
 71:   PetscCall(DMCreateMatrix(dm, pA));
 72:   A = *pA;
 73:   PetscCall(DMCreateGlobalVector(dm, pb));
 74:   b = *pb;
 75:   PetscCall(DMStagGetCorners(dm, &start, NULL, NULL, &n, NULL, NULL, &n_extra, NULL, NULL));
 76:   PetscCall(DMStagGetGlobalSizes(dm, &N, NULL, NULL));

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

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

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

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

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

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

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

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

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

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

135:   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
136:   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
137:   PetscCall(VecAssemblyBegin(b));
138:   PetscCall(VecAssemblyEnd(b));
139:   PetscFunctionReturn(PETSC_SUCCESS);
140: }

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

146:   PetscFunctionBeginUser;
147:   PetscCall(DMSetMatrixPreallocateOnly(dm, PETSC_TRUE));
148:   PetscCall(DMGetDimension(dm, &dim));
149:   switch (dim) {
150:   case 1:
151:     PetscCall(AssembleSystem1DVertexCentered(dm, pA, pb));
152:     break;
153:   default:
154:     SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_SUP, "Unsupported dimension: %" PetscInt_FMT, dim);
155:   }
156:   PetscFunctionReturn(PETSC_SUCCESS);
157: }

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

163: TEST*/