Date: 2025-03-29T05:58:24+00:00

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