Actual source code: ex8.c

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