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:   *isplit = ts->tsrhssplit;
  8:   /* look up the split */
  9:   while (*isplit) {
 10:     PetscStrcmp((*isplit)->splitname,splitname,&found);
 11:     if (found) break;
 12:     *isplit = (*isplit)->next;
 13:   }
 14:   return 0;
 15: }

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

 20:    Logically Collective on TS

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

 27:    Level: intermediate

 29: .seealso: TSRHSSplitGetIS()

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


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

 51:   PetscObjectIncrementTabLevel((PetscObject)newsplit->ts,(PetscObject)ts,1);
 52:   PetscLogObjectParent((PetscObject)ts,(PetscObject)newsplit->ts);
 53:   PetscSNPrintf(prefix,sizeof(prefix),"%srhsplit_%s_",((PetscObject)ts)->prefix ? ((PetscObject)ts)->prefix : "",newsplit->splitname);
 54:   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:   return 0;
 62: }

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

 67:    Logically Collective on TS

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

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

 76:    Level: intermediate

 78: .seealso: TSRHSSplitSetIS()

 80: @*/
 81: PetscErrorCode TSRHSSplitGetIS(TS ts,const char splitname[],IS *is)
 82: {
 83:   TS_RHSSplitLink isplit;

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

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

 96:    Logically Collective on TS

 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:  Calling sequence of fun:
106: $  rhsfunc(TS ts,PetscReal t,Vec u,Vec f,ctx);

108: +  t    - time at step/stage being solved
109: .  u    - state vector
110: .  f    - function vector
111: -  ctx  - [optional] user-defined context for matrix evaluation routine (may be NULL)

113:  Level: beginner

115: @*/
116: PetscErrorCode TSRHSSplitSetRHSFunction(TS ts,const char splitname[],Vec r,TSRHSFunction rhsfunc,void *ctx)
117: {
118:   TS_RHSSplitLink isplit;
119:   DM              dmc;
120:   Vec             subvec,ralloc = NULL;


125:   /* look up the split */
126:   TSRHSSplitGetRHSSplit(ts,splitname,&isplit);

129:   if (!r && ts->vec_sol) {
130:     VecGetSubVector(ts->vec_sol,isplit->is,&subvec);
131:     VecDuplicate(subvec,&ralloc);
132:     r    = ralloc;
133:     VecRestoreSubVector(ts->vec_sol,isplit->is,&subvec);
134:   }

136:   if (ts->dm) {
137:     PetscInt dim;

139:     DMGetDimension(ts->dm, &dim);
140:     if (dim != -1) {
141:       DMClone(ts->dm, &dmc);
142:       TSSetDM(isplit->ts, dmc);
143:       DMDestroy(&dmc);
144:     }
145:   }

147:   TSSetRHSFunction(isplit->ts,r,rhsfunc,ctx);
148:   VecDestroy(&ralloc);
149:   return 0;
150: }

152: /*@C
153:    TSRHSSplitGetSubTS - Get the sub-TS by split name.

155:    Logically Collective on TS

157:    Input Parameter:
158: .  ts - the TS context obtained from TSCreate()

160:    Output Parameters:
161: +  splitname - the number of the split
162: -  subts - the array of TS contexts

164:    Level: advanced

166: .seealso: TSGetRHSSplitFunction()
167: @*/
168: PetscErrorCode TSRHSSplitGetSubTS(TS ts,const char splitname[],TS *subts)
169: {
170:   TS_RHSSplitLink isplit;

174:   *subts = NULL;
175:   /* look up the split */
176:   TSRHSSplitGetRHSSplit(ts,splitname,&isplit);
177:   if (isplit) *subts = isplit->ts;
178:   return 0;
179: }

181: /*@C
182:    TSRHSSplitGetSubTSs - Get an array of all sub-TS contexts.

184:    Logically Collective on TS

186:    Input Parameter:
187: .  ts - the TS context obtained from TSCreate()

189:    Output Parameters:
190: +  n - the number of splits
191: -  subksp - the array of TS contexts

193:    Note:
194:    After TSRHSSplitGetSubTS() the array of TSs is to be freed by the user with PetscFree()
195:    (not the TS just the array that contains them).

197:    Level: advanced

199: .seealso: TSGetRHSSplitFunction()
200: @*/
201: PetscErrorCode TSRHSSplitGetSubTSs(TS ts,PetscInt *n,TS *subts[])
202: {
203:   TS_RHSSplitLink ilink = ts->tsrhssplit;
204:   PetscInt        i = 0;

207:   if (subts) {
208:     PetscMalloc1(ts->num_rhs_splits,subts);
209:     while (ilink) {
210:       (*subts)[i++] = ilink->ts;
211:       ilink = ilink->next;
212:     }
213:   }
214:   if (n) *n = ts->num_rhs_splits;
215:   return 0;
216: }