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: }