Actual source code: tsrhssplit.c

      1: #include <petsc/private/tsimpl.h>
      2: #include <petscdm.h>
      3: static PetscErrorCode TSRHSSplitGetRHSSplit(TS ts, const char splitname[], TS_RHSSplitLink *isplit)
      4: {
      5:   PetscBool found = PETSC_FALSE;

      7:   PetscFunctionBegin;
      8:   *isplit = ts->tsrhssplit;
      9:   /* look up the split */
     10:   while (*isplit) {
     11:     PetscCall(PetscStrcmp((*isplit)->splitname, splitname, &found));
     12:     if (found) break;
     13:     *isplit = (*isplit)->next;
     14:   }
     15:   PetscFunctionReturn(PETSC_SUCCESS);
     16: }

     18: /*@
     19:   TSRHSSplitSetIS - Set the index set for the specified split

     21:   Logically Collective

     23:   Input Parameters:
     24: + ts        - the `TS` context obtained from `TSCreate()`
     25: . splitname - name of this split, if `NULL` the number of the split is used
     26: - is        - the index set for part of the solution vector

     28:   Level: intermediate

     30: .seealso: [](ch_ts), `TS`, `IS`, `TSRHSSplitGetIS()`
     31: @*/
     32: PetscErrorCode TSRHSSplitSetIS(TS ts, const char splitname[], IS is)
     33: {
     34:   TS_RHSSplitLink newsplit, next = ts->tsrhssplit;
     35:   char            prefix[128];

     37:   PetscFunctionBegin;

     41:   PetscCall(PetscNew(&newsplit));
     42:   if (splitname) {
     43:     PetscCall(PetscStrallocpy(splitname, &newsplit->splitname));
     44:   } else {
     45:     PetscCall(PetscMalloc1(8, &newsplit->splitname));
     46:     PetscCall(PetscSNPrintf(newsplit->splitname, 7, "%" PetscInt_FMT, ts->num_rhs_splits));
     47:   }
     48:   PetscCall(PetscObjectReference((PetscObject)is));
     49:   newsplit->is = is;
     50:   PetscCall(TSCreate(PetscObjectComm((PetscObject)ts), &newsplit->ts));

     52:   PetscCall(PetscObjectIncrementTabLevel((PetscObject)newsplit->ts, (PetscObject)ts, 1));
     53:   PetscCall(PetscSNPrintf(prefix, sizeof(prefix), "%srhsplit_%s_", ((PetscObject)ts)->prefix ? ((PetscObject)ts)->prefix : "", newsplit->splitname));
     54:   PetscCall(TSSetOptionsPrefix(newsplit->ts, prefix));
     55:   if (!next) ts->tsrhssplit = newsplit;
     56:   else {
     57:     while (next->next) next = next->next;
     58:     next->next = newsplit;
     59:   }
     60:   ts->num_rhs_splits++;
     61:   PetscFunctionReturn(PETSC_SUCCESS);
     62: }

     64: /*@
     65:   TSRHSSplitGetIS - Retrieves the elements for a split as an `IS`

     67:   Logically Collective

     69:   Input Parameters:
     70: + ts        - the `TS` context obtained from `TSCreate()`
     71: - splitname - name of this split

     73:   Output Parameter:
     74: . is - the index set for part of the solution vector

     76:   Level: intermediate

     78: .seealso: [](ch_ts), `TS`, `IS`, `TSRHSSplitSetIS()`
     79: @*/
     80: PetscErrorCode TSRHSSplitGetIS(TS ts, const char splitname[], IS *is)
     81: {
     82:   TS_RHSSplitLink isplit;

     84:   PetscFunctionBegin;
     86:   *is = NULL;
     87:   /* look up the split */
     88:   PetscCall(TSRHSSplitGetRHSSplit(ts, splitname, &isplit));
     89:   if (isplit) *is = isplit->is;
     90:   PetscFunctionReturn(PETSC_SUCCESS);
     91: }

     93: /*@C
     94:   TSRHSSplitSetRHSFunction - Set the split right-hand-side functions.

     96:   Logically Collective

     98:   Input Parameters:
     99: + ts        - the `TS` context obtained from `TSCreate()`
    100: . splitname - name of this split
    101: . r         - vector to hold the residual (or `NULL` to have it created internally)
    102: . rhsfunc   - the RHS function evaluation routine
    103: - ctx       - user-defined context for private data for the split function evaluation routine (may be `NULL`)

    105:   Level: intermediate

    107: .seealso: [](ch_ts), `TS`, `TSRHSFunctionFn`, `IS`, `TSRHSSplitSetIS()`
    108: @*/
    109: PetscErrorCode TSRHSSplitSetRHSFunction(TS ts, const char splitname[], Vec r, TSRHSFunctionFn *rhsfunc, void *ctx)
    110: {
    111:   TS_RHSSplitLink isplit;
    112:   DM              dmc;
    113:   Vec             subvec, ralloc = NULL;

    115:   PetscFunctionBegin;

    119:   /* look up the split */
    120:   PetscCall(TSRHSSplitGetRHSSplit(ts, splitname, &isplit));
    121:   PetscCheck(isplit, PETSC_COMM_SELF, PETSC_ERR_USER, "The split %s is not created, check the split name or call TSRHSSplitSetIS() to create one", splitname);

    123:   if (!r && ts->vec_sol) {
    124:     PetscCall(VecGetSubVector(ts->vec_sol, isplit->is, &subvec));
    125:     PetscCall(VecDuplicate(subvec, &ralloc));
    126:     r = ralloc;
    127:     PetscCall(VecRestoreSubVector(ts->vec_sol, isplit->is, &subvec));
    128:   }

    130:   if (ts->dm) {
    131:     PetscInt dim;

    133:     PetscCall(DMGetDimension(ts->dm, &dim));
    134:     if (dim != -1) {
    135:       PetscCall(DMClone(ts->dm, &dmc));
    136:       PetscCall(TSSetDM(isplit->ts, dmc));
    137:       PetscCall(DMDestroy(&dmc));
    138:     }
    139:   }

    141:   PetscCall(TSSetRHSFunction(isplit->ts, r, rhsfunc, ctx));
    142:   PetscCall(VecDestroy(&ralloc));
    143:   PetscFunctionReturn(PETSC_SUCCESS);
    144: }

    146: /*@C
    147:   TSRHSSplitSetIFunction - Set the split implicit function for `TSARKIMEX`

    149:   Logically Collective

    151:   Input Parameters:
    152: + ts        - the `TS` context obtained from `TSCreate()`
    153: . splitname - name of this split
    154: . r         - vector to hold the residual (or `NULL` to have it created internally)
    155: . ifunc     - the implicit function evaluation routine
    156: - ctx       - user-defined context for private data for the split function evaluation routine (may be `NULL`)

    158:   Level: intermediate

    160: .seealso: [](ch_ts), `TS`, `TSIFunctionFn`, `IS`, `TSRHSSplitSetIS()`, `TSARKIMEX`
    161: @*/
    162: PetscErrorCode TSRHSSplitSetIFunction(TS ts, const char splitname[], Vec r, TSIFunctionFn *ifunc, void *ctx)
    163: {
    164:   TS_RHSSplitLink isplit;
    165:   DM              dmc;
    166:   Vec             subvec, ralloc = NULL;

    168:   PetscFunctionBegin;

    172:   /* look up the split */
    173:   PetscCall(TSRHSSplitGetRHSSplit(ts, splitname, &isplit));
    174:   PetscCheck(isplit, PETSC_COMM_SELF, PETSC_ERR_USER, "The split %s is not created, check the split name or call TSRHSSplitSetIS() to create one", splitname);

    176:   if (!r && ts->vec_sol) {
    177:     PetscCall(VecGetSubVector(ts->vec_sol, isplit->is, &subvec));
    178:     PetscCall(VecDuplicate(subvec, &ralloc));
    179:     r = ralloc;
    180:     PetscCall(VecRestoreSubVector(ts->vec_sol, isplit->is, &subvec));
    181:   }

    183:   if (ts->dm) {
    184:     PetscInt dim;

    186:     PetscCall(DMGetDimension(ts->dm, &dim));
    187:     if (dim != -1) {
    188:       PetscCall(DMClone(ts->dm, &dmc));
    189:       PetscCall(TSSetDM(isplit->ts, dmc));
    190:       PetscCall(DMDestroy(&dmc));
    191:     }
    192:   }

    194:   PetscCall(TSSetIFunction(isplit->ts, r, ifunc, ctx));
    195:   PetscCall(VecDestroy(&ralloc));
    196:   PetscFunctionReturn(PETSC_SUCCESS);
    197: }

    199: /*@C
    200:   TSRHSSplitSetIJacobian - Set the Jacobian for the split implicit function with `TSARKIMEX`

    202:   Logically Collective

    204:   Input Parameters:
    205: + ts        - the `TS` context obtained from `TSCreate()`
    206: . splitname - name of this split
    207: . Amat      - (approximate) matrix to store Jacobian entries computed by `f`
    208: . Pmat      - matrix used to compute preconditioner (usually the same as `Amat`)
    209: . ijac      - the Jacobian evaluation routine
    210: - ctx       - user-defined context for private data for the split function evaluation routine (may be `NULL`)

    212:   Level: intermediate

    214: .seealso: [](ch_ts), `TS`, `TSRHSSplitSetIFunction`, `TSIJacobianFn`, `IS`, `TSRHSSplitSetIS()`
    215: @*/
    216: PetscErrorCode TSRHSSplitSetIJacobian(TS ts, const char splitname[], Mat Amat, Mat Pmat, TSIJacobianFn *ijac, void *ctx)
    217: {
    218:   TS_RHSSplitLink isplit;
    219:   DM              dmc;

    221:   PetscFunctionBegin;
    225:   if (Amat) PetscCheckSameComm(ts, 1, Amat, 3);
    226:   if (Pmat) PetscCheckSameComm(ts, 1, Pmat, 4);

    228:   /* look up the split */
    229:   PetscCall(TSRHSSplitGetRHSSplit(ts, splitname, &isplit));
    230:   PetscCheck(isplit, PETSC_COMM_SELF, PETSC_ERR_USER, "The split %s is not created, check the split name or call TSRHSSplitSetIS() to create one", splitname);

    232:   if (ts->dm) {
    233:     PetscInt dim;

    235:     PetscCall(DMGetDimension(ts->dm, &dim));
    236:     if (dim != -1) {
    237:       PetscCall(DMClone(ts->dm, &dmc));
    238:       PetscCall(TSSetDM(isplit->ts, dmc));
    239:       PetscCall(DMDestroy(&dmc));
    240:     }
    241:   }

    243:   PetscCall(TSSetIJacobian(isplit->ts, Amat, Pmat, ijac, ctx));
    244:   PetscFunctionReturn(PETSC_SUCCESS);
    245: }

    247: /*@C
    248:   TSRHSSplitGetSubTS - Get the sub-`TS` by split name.

    250:   Logically Collective

    252:   Input Parameter:
    253: . ts - the `TS` context obtained from `TSCreate()`

    255:   Output Parameters:
    256: + splitname - the number of the split
    257: - subts     - the sub-`TS`

    259:   Level: advanced

    261: .seealso: [](ch_ts), `TS`, `IS`, `TSGetRHSSplitFunction()`
    262: @*/
    263: PetscErrorCode TSRHSSplitGetSubTS(TS ts, const char splitname[], TS *subts)
    264: {
    265:   TS_RHSSplitLink isplit;

    267:   PetscFunctionBegin;
    269:   PetscAssertPointer(subts, 3);
    270:   *subts = NULL;
    271:   /* look up the split */
    272:   PetscCall(TSRHSSplitGetRHSSplit(ts, splitname, &isplit));
    273:   if (isplit) *subts = isplit->ts;
    274:   PetscFunctionReturn(PETSC_SUCCESS);
    275: }

    277: /*@C
    278:   TSRHSSplitGetSubTSs - Get an array of all sub-`TS` contexts.

    280:   Logically Collective

    282:   Input Parameter:
    283: . ts - the `TS` context obtained from `TSCreate()`

    285:   Output Parameters:
    286: + n     - the number of splits
    287: - subts - the array of `TS` contexts

    289:   Level: advanced

    291:   Note:
    292:   After `TSRHSSplitGetSubTS()` the array of `TS`s is to be freed by the user with `PetscFree()`
    293:   (not the `TS` in the array just the array that contains them).

    295: .seealso: [](ch_ts), `TS`, `IS`, `TSGetRHSSplitFunction()`
    296: @*/
    297: PetscErrorCode TSRHSSplitGetSubTSs(TS ts, PetscInt *n, TS *subts[])
    298: {
    299:   TS_RHSSplitLink ilink = ts->tsrhssplit;
    300:   PetscInt        i     = 0;

    302:   PetscFunctionBegin;
    304:   if (subts) {
    305:     PetscCall(PetscMalloc1(ts->num_rhs_splits, subts));
    306:     while (ilink) {
    307:       (*subts)[i++] = ilink->ts;
    308:       ilink         = ilink->next;
    309:     }
    310:   }
    311:   if (n) *n = ts->num_rhs_splits;
    312:   PetscFunctionReturn(PETSC_SUCCESS);
    313: }

    315: /*@
    316:   TSRHSSplitGetSNES - Returns the `SNES` (nonlinear solver) associated with
    317:   a `TS` (timestepper) context when RHS splits are used.

    319:   Not Collective, but snes is parallel if ts is parallel

    321:   Input Parameter:
    322: . ts - the `TS` context obtained from `TSCreate()`

    324:   Output Parameter:
    325: . snes - the nonlinear solver context

    327:   Level: intermediate

    329:   Note:
    330:   The returned `SNES` may have a different `DM` with the `TS` `DM`.

    332: .seealso: [](ch_ts), `TS`, `SNES`, `TSCreate()`, `TSRHSSplitSetSNES()`
    333: @*/
    334: PetscErrorCode TSRHSSplitGetSNES(TS ts, SNES *snes)
    335: {
    336:   PetscFunctionBegin;
    338:   PetscAssertPointer(snes, 2);
    339:   if (!ts->snesrhssplit) {
    340:     PetscCall(SNESCreate(PetscObjectComm((PetscObject)ts), &ts->snesrhssplit));
    341:     PetscCall(PetscObjectSetOptions((PetscObject)ts->snesrhssplit, ((PetscObject)ts)->options));
    342:     PetscCall(SNESSetFunction(ts->snesrhssplit, NULL, SNESTSFormFunction, ts));
    343:     PetscCall(PetscObjectIncrementTabLevel((PetscObject)ts->snesrhssplit, (PetscObject)ts, 1));
    344:     if (ts->problem_type == TS_LINEAR) PetscCall(SNESSetType(ts->snesrhssplit, SNESKSPONLY));
    345:   }
    346:   *snes = ts->snesrhssplit;
    347:   PetscFunctionReturn(PETSC_SUCCESS);
    348: }

    350: /*@
    351:   TSRHSSplitSetSNES - Set the `SNES` (nonlinear solver) to be used by the
    352:   timestepping context when RHS splits are used.

    354:   Collective

    356:   Input Parameters:
    357: + ts   - the `TS` context obtained from `TSCreate()`
    358: - snes - the nonlinear solver context

    360:   Level: intermediate

    362:   Note:
    363:   Most users should have the `TS` created by calling `TSRHSSplitGetSNES()`

    365: .seealso: [](ch_ts), `TS`, `SNES`, `TSCreate()`, `TSRHSSplitGetSNES()`
    366: @*/
    367: PetscErrorCode TSRHSSplitSetSNES(TS ts, SNES snes)
    368: {
    369:   PetscErrorCode (*func)(SNES, Vec, Mat, Mat, void *);

    371:   PetscFunctionBegin;
    374:   PetscCall(PetscObjectReference((PetscObject)snes));
    375:   PetscCall(SNESDestroy(&ts->snesrhssplit));

    377:   ts->snesrhssplit = snes;

    379:   PetscCall(SNESSetFunction(ts->snesrhssplit, NULL, SNESTSFormFunction, ts));
    380:   PetscCall(SNESGetJacobian(ts->snesrhssplit, NULL, NULL, &func, NULL));
    381:   if (func == SNESTSFormJacobian) PetscCall(SNESSetJacobian(ts->snesrhssplit, NULL, NULL, SNESTSFormJacobian, ts));
    382:   PetscFunctionReturn(PETSC_SUCCESS);
    383: }