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: /*@C
 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: /*@C
 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:   TSRHSSplitGetSubTS - Get the sub-`TS` by split name.

149:   Logically Collective

151:   Input Parameter:
152: . ts - the `TS` context obtained from `TSCreate()`

154:   Output Parameters:
155: + splitname - the number of the split
156: - subts     - the sub-`TS`

158:   Level: advanced

160: .seealso: [](ch_ts), `TS`, `IS`, `TSGetRHSSplitFunction()`
161: @*/
162: PetscErrorCode TSRHSSplitGetSubTS(TS ts, const char splitname[], TS *subts)
163: {
164:   TS_RHSSplitLink isplit;

166:   PetscFunctionBegin;
168:   PetscAssertPointer(subts, 3);
169:   *subts = NULL;
170:   /* look up the split */
171:   PetscCall(TSRHSSplitGetRHSSplit(ts, splitname, &isplit));
172:   if (isplit) *subts = isplit->ts;
173:   PetscFunctionReturn(PETSC_SUCCESS);
174: }

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

179:   Logically Collective

181:   Input Parameter:
182: . ts - the `TS` context obtained from `TSCreate()`

184:   Output Parameters:
185: + n     - the number of splits
186: - subts - the array of `TS` contexts

188:   Level: advanced

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

194: .seealso: [](ch_ts), `TS`, `IS`, `TSGetRHSSplitFunction()`
195: @*/
196: PetscErrorCode TSRHSSplitGetSubTSs(TS ts, PetscInt *n, TS *subts[])
197: {
198:   TS_RHSSplitLink ilink = ts->tsrhssplit;
199:   PetscInt        i     = 0;

201:   PetscFunctionBegin;
203:   if (subts) {
204:     PetscCall(PetscMalloc1(ts->num_rhs_splits, subts));
205:     while (ilink) {
206:       (*subts)[i++] = ilink->ts;
207:       ilink         = ilink->next;
208:     }
209:   }
210:   if (n) *n = ts->num_rhs_splits;
211:   PetscFunctionReturn(PETSC_SUCCESS);
212: }