Date: 2025-03-28T21:10:49+00:00

Actual source code: ex1.c

      1: static char help[] = "Solve a toy 1D problem on a staggered grid.\n\
      2:                       Accepts command line options -a, -b, and -c\n\
      3:                       and approximately solves\n\
      4:                       u''(x) = c, u(0) = 1, u(1) = b\n\n";
      5: /*

      7:    To demonstrate the basic functionality of DMStag, solves a second-order ODE,

      9:        u''(x) = f(x),  0 < x < 1
     10:        u(0) = a
     11:        u(1) = b

     13:    in mixed form, that is by introducing an auxiliary variable p

     15:       p'(x) = f(x), p - u'(x) = 0, 0 < x < 1
     16:       u(0) = a
     17:       u(1) = b

     19:    For f == c, the solution is

     21:      u(x) = a + (b - a - (c/2)) x + (c/2) x^2
     22:      p(x) = (b - a - (c/2)) + c x

     24:    To find an approximate solution, discretize by storing values of p in
     25:    elements and values of u on their boundaries, and using first-order finite
     26:    differences.

     28:    This should in fact produce a (nodal) solution with no discretization error,
     29:    so differences from the reference solution will be restricted to those induced
     30:    by floating point operations (in particular, the finite differences) and the
     31:    accuracy of the linear solve.

     33:    Parameters for the main grid can be controlled with command line options, e.g.

     35:      -stag_grid_x 10

     37:   In particular to notice in this example are the two methods of indexing. The
     38:   first is analogous to the use of MatStencil with DMDA, and the second is
     39:   analogous to the use of DMDAVecGetArrayDOF().

     41:   The first, recommended for ease of use, is based on naming an element in the
     42:   global grid, a location in its support, and a component. For example,
     43:   one might consider element e, the left side (a vertex in 1d), and the first
     44:   component (index 0). This is accomplished by populating a DMStagStencil struct,
     45:   e.g.

     47:       DMStagStencil stencil;
     48:       stencil.i   = i
     49:       stencil.loc = DMSTAG_LEFT;
     50:       stencil.c   = 0

     52:   Note that below, for convenenience, we #define an alias LEFT for DMSTAG_LEFT.

     54:   The second, which ensures maximum efficiency, makes use of the underlying
     55:   block structure of local DMStag-derived vectors, and requires the user to
     56:   obtain the correct local offset for the degrees of freedom they would like to
     57:   use. This is made simple with the helper function DMStagGetLocationSlot().

     59:   Note that the linear system being solved is indefinite, so is not entirely
     60:   trivial to invert. The default solver here (GMRES/Jacobi) is a poor choice,
     61:   made to avoid depending on an external package. To solve a larger system,
     62:   the usual method for a 1-d problem such as this is to choose a sophisticated
     63:   direct solver, e.g. configure --download-suitesparse and run

     65:     $PETSC_DIR/$PETSC_ARCH/bin/mpiexec -n 3 ./stag_ex2 -stag_grid_x 100 -pc_type lu -pc_factor_mat_solver_package umfpack

     67:   You can also impose a periodic boundary condition, in which case -b and -c are
     68:   ignored; b = a and c = 0.0 are used, giving a constant u == a , p == 0.

     70:       -stag_boundary_type_x periodic

     72: */
     73: #include <petscdm.h>
     74: #include <petscksp.h>
     75: #include <petscdmstag.h>

     77: /* Shorter, more convenient names for DMStagStencilLocation entries */
     78: #define LEFT    DMSTAG_LEFT
     79: #define RIGHT   DMSTAG_RIGHT
     80: #define ELEMENT DMSTAG_ELEMENT

     82: int main(int argc, char **argv)
     83: {
     84:   DM             dmSol, dmForcing;
     85:   DM             dmCoordSol;
     86:   Vec            sol, solRef, solRefLocal, f, fLocal, rhs, coordSolLocal;
     87:   Mat            A;
     88:   PetscScalar    a, b, c, h;
     89:   KSP            ksp;
     90:   PC             pc;
     91:   PetscInt       start, n, e, nExtra;
     92:   PetscInt       iu, ip, ixu, ixp;
     93:   PetscBool      isLastRank, isFirstRank;
     94:   PetscScalar  **arrSol, **arrCoordSol;
     95:   DMBoundaryType boundary;

     97:   const PetscReal domainSize = 1.0;

     99:   /* Initialize PETSc */
    100:   PetscFunctionBeginUser;
    101:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));

    103:   /* Create 1D DMStag for the solution, and set up. Note that you can supply many
    104:      command line options (see the man page for DMStagCreate1d)
    105:   */
    106:   PetscCall(DMStagCreate1d(PETSC_COMM_WORLD, DM_BOUNDARY_NONE, 3, 1, 1, DMSTAG_STENCIL_BOX, 1, NULL, &dmSol));
    107:   PetscCall(DMSetFromOptions(dmSol));
    108:   PetscCall(DMSetUp(dmSol));

    110:   /* Create uniform coordinates. Note that in higher-dimensional examples,
    111:       coordinates are created differently.*/
    112:   PetscCall(DMStagSetUniformCoordinatesExplicit(dmSol, 0.0, domainSize, 0.0, 0.0, 0.0, 0.0));

    114:   /* Determine boundary type */
    115:   PetscCall(DMStagGetBoundaryTypes(dmSol, &boundary, NULL, NULL));

    117:   /* Process command line options (depends on DMStag setup) */
    118:   a = 1.0;
    119:   b = 2.0;
    120:   c = 1.0;
    121:   PetscCall(PetscOptionsGetScalar(NULL, NULL, "-a", &a, NULL));
    122:   if (boundary == DM_BOUNDARY_PERIODIC) {
    123:     b = a;
    124:     c = 0.0;
    125:   } else {
    126:     PetscCall(PetscOptionsGetScalar(NULL, NULL, "-b", &b, NULL));
    127:     PetscCall(PetscOptionsGetScalar(NULL, NULL, "-c", &c, NULL));
    128:   }

    130:   /* Compute reference solution on the grid, using direct array access */
    131:   PetscCall(DMCreateGlobalVector(dmSol, &solRef));
    132:   PetscCall(DMGetLocalVector(dmSol, &solRefLocal));
    133:   PetscCall(DMStagVecGetArray(dmSol, solRefLocal, &arrSol));
    134:   PetscCall(DMGetCoordinateDM(dmSol, &dmCoordSol));
    135:   PetscCall(DMGetCoordinatesLocal(dmSol, &coordSolLocal));
    136:   PetscCall(DMStagVecGetArrayRead(dmCoordSol, coordSolLocal, &arrCoordSol));
    137:   PetscCall(DMStagGetCorners(dmSol, &start, NULL, NULL, &n, NULL, NULL, &nExtra, NULL, NULL));

    139:   /* Get the correct entries for each of our variables in local element-wise storage */
    140:   PetscCall(DMStagGetLocationSlot(dmSol, LEFT, 0, &iu));
    141:   PetscCall(DMStagGetLocationSlot(dmSol, ELEMENT, 0, &ip));
    142:   PetscCall(DMStagGetLocationSlot(dmCoordSol, LEFT, 0, &ixu));
    143:   PetscCall(DMStagGetLocationSlot(dmCoordSol, ELEMENT, 0, &ixp));
    144:   for (e = start; e < start + n + nExtra; ++e) {
    145:     {
    146:       const PetscScalar coordu = arrCoordSol[e][ixu];
    147:       arrSol[e][iu]            = a + (b - a - (c / 2.0)) * coordu + (c / 2.0) * coordu * coordu;
    148:     }
    149:     if (e < start + n) {
    150:       const PetscScalar coordp = arrCoordSol[e][ixp];
    151:       arrSol[e][ip]            = b - a - (c / 2.0) + c * coordp;
    152:     }
    153:   }
    154:   PetscCall(DMStagVecRestoreArrayRead(dmCoordSol, coordSolLocal, &arrCoordSol));
    155:   PetscCall(DMStagVecRestoreArray(dmSol, solRefLocal, &arrSol));
    156:   PetscCall(DMLocalToGlobal(dmSol, solRefLocal, INSERT_VALUES, solRef));
    157:   PetscCall(DMRestoreLocalVector(dmSol, &solRefLocal));

    159:   /* Create another 1D DMStag for the forcing term, and populate a field on it.
    160:      Here this is not really necessary, but in other contexts we may have auxiliary
    161:      fields which we use to construct the linear system.

    163:      This second DM represents the same physical domain, but has a different
    164:      "default section" (though the current implementation does NOT actually use
    165:      PetscSection). Since it is created as a derivative of the original DMStag,
    166:      we can be confident that it is compatible. One could check with DMGetCompatiblity() */
    167:   PetscCall(DMStagCreateCompatibleDMStag(dmSol, 1, 0, 0, 0, &dmForcing));
    168:   PetscCall(DMCreateGlobalVector(dmForcing, &f));
    169:   PetscCall(VecSet(f, c)); /* Dummy for logic which depends on auxiliary data */

    171:   /* Assemble System */
    172:   PetscCall(DMCreateMatrix(dmSol, &A));
    173:   PetscCall(DMCreateGlobalVector(dmSol, &rhs));
    174:   PetscCall(DMGetLocalVector(dmForcing, &fLocal));
    175:   PetscCall(DMGlobalToLocal(dmForcing, f, INSERT_VALUES, fLocal));

    177:   /* Note: if iterating over all the elements, you will usually need to do something
    178:      special at one of the boundaries. You can either make use of the existence
    179:      of a "extra" partial dummy element on the right/top/front, or you can use a different stencil.
    180:      The construction of the reference solution above uses the first method,
    181:      so here we will use the second */

    183:   PetscCall(DMStagGetIsLastRank(dmSol, &isLastRank, NULL, NULL));
    184:   PetscCall(DMStagGetIsFirstRank(dmSol, &isFirstRank, NULL, NULL));
    185:   for (e = start; e < start + n; ++e) {
    186:     DMStagStencil pos[3];
    187:     PetscScalar   val[3];
    188:     PetscInt      idxLoc;

    190:     idxLoc          = 0;
    191:     pos[idxLoc].i   = e;       /* This element in the 1d ordering */
    192:     pos[idxLoc].loc = ELEMENT; /* Element-centered dofs (p) */
    193:     pos[idxLoc].c   = 0;       /* Component 0 : first (and only) p dof */
    194:     val[idxLoc]     = 0.0;     /* p - u'(x) = 0 */
    195:     ++idxLoc;

    197:     if (isFirstRank && e == start) {
    198:       /* Special case on left boundary */
    199:       pos[idxLoc].i   = e;    /* This element in the 1d ordering */
    200:       pos[idxLoc].loc = LEFT; /* Left vertex */
    201:       pos[idxLoc].c   = 0;
    202:       val[idxLoc]     = a; /* u(0) = a */
    203:       ++idxLoc;
    204:     } else {
    205:       PetscScalar fVal;
    206:       /* Usual case - deal with velocity on left side of cell
    207:          Here, we obtain a value of f (even though it's constant here,
    208:          this demonstrates the more-realistic case of a pre-computed coefficient) */
    209:       pos[idxLoc].i   = e;    /* This element in the 1d ordering */
    210:       pos[idxLoc].loc = LEFT; /* vertex-centered dof (u) */
    211:       pos[idxLoc].c   = 0;

    213:       PetscCall(DMStagVecGetValuesStencil(dmForcing, fLocal, 1, &pos[idxLoc], &fVal));

    215:       val[idxLoc] = fVal; /* p'(x) = f, in interior */
    216:       ++idxLoc;
    217:     }
    218:     if (boundary != DM_BOUNDARY_PERIODIC && isLastRank && e == start + n - 1) {
    219:       /* Special case on right boundary (in addition to usual case) */
    220:       pos[idxLoc].i   = e; /* This element in the 1d ordering */
    221:       pos[idxLoc].loc = RIGHT;
    222:       pos[idxLoc].c   = 0;
    223:       val[idxLoc]     = b; /* u(1) = b */
    224:       ++idxLoc;
    225:     }
    226:     PetscCall(DMStagVecSetValuesStencil(dmSol, rhs, idxLoc, pos, val, INSERT_VALUES));
    227:   }
    228:   PetscCall(DMRestoreLocalVector(dmForcing, &fLocal));
    229:   PetscCall(VecAssemblyBegin(rhs));
    230:   PetscCall(VecAssemblyEnd(rhs));

    232:   /* Note: normally it would be more efficient to assemble the RHS and the matrix
    233:      in the same loop over elements, but we separate them for clarity here */
    234:   PetscCall(DMGetCoordinatesLocal(dmSol, &coordSolLocal));
    235:   for (e = start; e < start + n; ++e) {
    236:     /* Velocity is either a BC or an interior point */
    237:     if (isFirstRank && e == start) {
    238:       DMStagStencil row;
    239:       PetscScalar   val;

    241:       row.i   = e;
    242:       row.loc = LEFT;
    243:       row.c   = 0;
    244:       val     = 1.0;
    245:       PetscCall(DMStagMatSetValuesStencil(dmSol, A, 1, &row, 1, &row, &val, INSERT_VALUES));
    246:     } else {
    247:       DMStagStencil row, col[3];
    248:       PetscScalar   val[3], xp[2];

    250:       row.i   = e;
    251:       row.loc = LEFT; /* In general, opt for LEFT/DOWN/BACK  and iterate over elements */
    252:       row.c   = 0;

    254:       col[0].i   = e;
    255:       col[0].loc = ELEMENT;
    256:       col[0].c   = 0;

    258:       col[1].i   = e - 1;
    259:       col[1].loc = ELEMENT;
    260:       col[1].c   = 0;

    262:       PetscCall(DMStagVecGetValuesStencil(dmCoordSol, coordSolLocal, 2, col, xp));
    263:       h = xp[0] - xp[1];
    264:       if (boundary == DM_BOUNDARY_PERIODIC && PetscRealPart(h) < 0.0) h += domainSize;

    266:       val[0] = 1.0 / h;
    267:       val[1] = -1.0 / h;

    269:       /* For convenience, we add an explicit 0 on the diagonal. This is a waste,
    270:          but it allows for easier use of a direct solver, if desired */
    271:       col[2].i   = e;
    272:       col[2].loc = LEFT;
    273:       col[2].c   = 0;
    274:       val[2]     = 0.0;

    276:       PetscCall(DMStagMatSetValuesStencil(dmSol, A, 1, &row, 3, col, val, INSERT_VALUES));
    277:     }

    279:     /* Additional velocity point (BC) on the right */
    280:     if (boundary != DM_BOUNDARY_PERIODIC && isLastRank && e == start + n - 1) {
    281:       DMStagStencil row;
    282:       PetscScalar   val;

    284:       row.i   = e;
    285:       row.loc = RIGHT;
    286:       row.c   = 0;
    287:       val     = 1.0;
    288:       PetscCall(DMStagMatSetValuesStencil(dmSol, A, 1, &row, 1, &row, &val, INSERT_VALUES));
    289:     }

    291:     /* Equation on pressure (element) variables */
    292:     {
    293:       DMStagStencil row, col[3];
    294:       PetscScalar   val[3], xu[2];

    296:       row.i   = e;
    297:       row.loc = ELEMENT;
    298:       row.c   = 0;

    300:       col[0].i   = e;
    301:       col[0].loc = RIGHT;
    302:       col[0].c   = 0;

    304:       col[1].i   = e;
    305:       col[1].loc = LEFT;
    306:       col[1].c   = 0;

    308:       PetscCall(DMStagVecGetValuesStencil(dmCoordSol, coordSolLocal, 2, col, xu));
    309:       h = xu[0] - xu[1];
    310:       if (boundary == DM_BOUNDARY_PERIODIC && PetscRealPart(h) < 0.0) h += domainSize;

    312:       val[0] = -1.0 / h;
    313:       val[1] = 1.0 / h;

    315:       col[2].i   = e;
    316:       col[2].loc = ELEMENT;
    317:       col[2].c   = 0;
    318:       val[2]     = 1.0;

    320:       PetscCall(DMStagMatSetValuesStencil(dmSol, A, 1, &row, 3, col, val, INSERT_VALUES));
    321:     }
    322:   }
    323:   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
    324:   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));

    326:   /* Solve */
    327:   PetscCall(DMCreateGlobalVector(dmSol, &sol));
    328:   PetscCall(KSPCreate(PETSC_COMM_WORLD, &ksp));
    329:   PetscCall(KSPSetOperators(ksp, A, A));
    330:   PetscCall(KSPGetPC(ksp, &pc));
    331:   PetscCall(PCSetType(pc, PCJACOBI)); /* A simple, but non-scalable, solver choice */
    332:   PetscCall(KSPSetFromOptions(ksp));
    333:   PetscCall(KSPSolve(ksp, rhs, sol));

    335:   /* View the components of the solution, demonstrating DMStagMigrateVec() */
    336:   {
    337:     DM  dmVertsOnly, dmElementsOnly;
    338:     Vec u, p;

    340:     PetscCall(DMStagCreateCompatibleDMStag(dmSol, 1, 0, 0, 0, &dmVertsOnly));
    341:     PetscCall(DMStagCreateCompatibleDMStag(dmSol, 0, 1, 0, 0, &dmElementsOnly));
    342:     PetscCall(DMGetGlobalVector(dmVertsOnly, &u));
    343:     PetscCall(DMGetGlobalVector(dmElementsOnly, &p));

    345:     PetscCall(DMStagMigrateVec(dmSol, sol, dmVertsOnly, u));
    346:     PetscCall(DMStagMigrateVec(dmSol, sol, dmElementsOnly, p));

    348:     PetscCall(PetscObjectSetName((PetscObject)u, "Sol_u"));
    349:     PetscCall(VecView(u, PETSC_VIEWER_STDOUT_WORLD));
    350:     PetscCall(PetscObjectSetName((PetscObject)p, "Sol_p"));
    351:     PetscCall(VecView(p, PETSC_VIEWER_STDOUT_WORLD));

    353:     PetscCall(DMRestoreGlobalVector(dmVertsOnly, &u));
    354:     PetscCall(DMRestoreGlobalVector(dmElementsOnly, &p));
    355:     PetscCall(DMDestroy(&dmVertsOnly));
    356:     PetscCall(DMDestroy(&dmElementsOnly));
    357:   }

    359:   /* Check Solution */
    360:   {
    361:     Vec       diff;
    362:     PetscReal normsolRef, errAbs, errRel;

    364:     PetscCall(VecDuplicate(sol, &diff));
    365:     PetscCall(VecCopy(sol, diff));
    366:     PetscCall(VecAXPY(diff, -1.0, solRef));
    367:     PetscCall(VecNorm(diff, NORM_2, &errAbs));
    368:     PetscCall(VecNorm(solRef, NORM_2, &normsolRef));
    369:     errRel = errAbs / normsolRef;
    370:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Error (abs): %g\nError (rel): %g\n", (double)errAbs, (double)errRel));
    371:     PetscCall(VecDestroy(&diff));
    372:   }

    374:   /* Clean up and finalize PETSc */
    375:   PetscCall(KSPDestroy(&ksp));
    376:   PetscCall(VecDestroy(&sol));
    377:   PetscCall(VecDestroy(&solRef));
    378:   PetscCall(VecDestroy(&rhs));
    379:   PetscCall(VecDestroy(&f));
    380:   PetscCall(MatDestroy(&A));
    381:   PetscCall(DMDestroy(&dmSol));
    382:   PetscCall(DMDestroy(&dmForcing));
    383:   PetscCall(PetscFinalize());
    384:   return 0;
    385: }

    387: /*TEST

    389:    test:
    390:       suffix: 1
    391:       nsize: 7
    392:       args: -dm_view -stag_grid_x 11 -stag_stencil_type star -a 1.33 -b 7.22 -c 347.2 -ksp_monitor_short

    394:    test:
    395:       suffix: periodic
    396:       nsize: 3
    397:       args: -dm_view -stag_grid_x 13 -stag_boundary_type_x periodic -a 1.1234

    399:    test:
    400:       suffix: periodic_seq
    401:       nsize: 1
    402:       args: -dm_view -stag_grid_x 13 -stag_boundary_type_x periodic -a 1.1234

    404:    test:
    405:       suffix: ghosted_vacuous
    406:       nsize: 3
    407:       args: -dm_view -stag_grid_x 13 -stag_boundary_type_x ghosted -a 1.1234

    409: TEST*/