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:   PetscCheck(dim == 1, PETSC_COMM_WORLD, PETSC_ERR_SUP, "Unsupported dimension: %" PetscInt_FMT, dim);
 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:   PetscCall(DMSetFromOptions(dm));
 32:   PetscCall(DMSetUp(dm));

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

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

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

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

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

 69:   PetscFunctionBeginUser;
 70:   PetscCall(DMCreateMatrix(dm, pA));
 71:   A = *pA;
 72:   PetscCall(DMCreateGlobalVector(dm, pb));
 73:   b = *pb;
 74:   PetscCall(DMStagGetCorners(dm, &start, NULL, NULL, &n, NULL, NULL, &n_extra, NULL, NULL));
 75:   PetscCall(DMStagGetGlobalSizes(dm, &N, NULL, NULL));

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

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

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

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

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

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

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

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

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

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

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

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

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

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

162: TEST*/