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;
  6:   PetscErrorCode  ierr;

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

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

 22:    Logically Collective on TS

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

 29:    Level: intermediate

 31: .seealso: TSRHSSplitGetIS()

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


 44:   PetscNew(&newsplit);
 45:   if (splitname) {
 46:     PetscStrallocpy(splitname,&newsplit->splitname);
 47:   } else {
 48:     PetscMalloc1(8,&newsplit->splitname);
 49:     PetscSNPrintf(newsplit->splitname,7,"%D",ts->num_rhs_splits);
 50:   }
 51:   PetscObjectReference((PetscObject)is);
 52:   newsplit->is = is;
 53:   TSCreate(PetscObjectComm((PetscObject)ts),&newsplit->ts);
 54:   PetscObjectIncrementTabLevel((PetscObject)newsplit->ts,(PetscObject)ts,1);
 55:   PetscLogObjectParent((PetscObject)ts,(PetscObject)newsplit->ts);
 56:   PetscSNPrintf(prefix,sizeof(prefix),"%srhsplit_%s_",((PetscObject)ts)->prefix ? ((PetscObject)ts)->prefix : "",newsplit->splitname);
 57:   TSSetOptionsPrefix(newsplit->ts,prefix);
 58:   if (!next) ts->tsrhssplit = newsplit;
 59:   else {
 60:     while (next->next) next = next->next;
 61:     next->next = newsplit;
 62:   }
 63:   ts->num_rhs_splits++;
 64:   return(0);
 65: }

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

 70:    Logically Collective on TS

 72:    Input Parameters:
 73: +  ts        - the TS context obtained from TSCreate()
 74: -  splitname - name of this split

 76:    Output Parameters:
 77: -  is        - the index set for part of the solution vector

 79:    Level: intermediate

 81: .seealso: TSRHSSplitSetIS()

 83: @*/
 84: PetscErrorCode TSRHSSplitGetIS(TS ts,const char splitname[],IS *is)
 85: {
 86:   TS_RHSSplitLink isplit;
 87:   PetscErrorCode  ierr;

 91:   *is = NULL;
 92:   /* look up the split */
 93:   TSRHSSplitGetRHSSplit(ts,splitname,&isplit);
 94:   if (isplit) *is = isplit->is;
 95:   return(0);
 96: }

 98: /*@C
 99:    TSRHSSplitSetRHSFunction - Set the split right-hand-side functions.

101:    Logically Collective on TS

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

110:  Calling sequence of fun:
111: $  rhsfunc(TS ts,PetscReal t,Vec u,Vec f,ctx);

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

118:  Level: beginner

120: @*/
121: PetscErrorCode TSRHSSplitSetRHSFunction(TS ts,const char splitname[],Vec r,TSRHSFunction rhsfunc,void *ctx)
122: {
123:   TS_RHSSplitLink isplit;
124:   DM              dmc;
125:   Vec             subvec,ralloc = NULL;
126:   PetscErrorCode  ierr;


132:   /* look up the split */
133:   TSRHSSplitGetRHSSplit(ts,splitname,&isplit);
134:   if (!isplit) SETERRQ1(PETSC_COMM_SELF,PETSC_ERR_USER,"The split %s is not created, check the split name or call TSRHSSplitSetIS() to create one",splitname);

136:   if (!r && ts->vec_sol) {
137:     VecGetSubVector(ts->vec_sol,isplit->is,&subvec);
138:     VecDuplicate(subvec,&ralloc);
139:     r    = ralloc;
140:     VecRestoreSubVector(ts->vec_sol,isplit->is,&subvec);
141:   }
142: 
143:   if(ts->dm){
144:     DMClone(ts->dm, &dmc);
145:     TSSetDM(isplit->ts, dmc);
146:     DMDestroy(&dmc);
147:   }
148: 
149:   TSSetRHSFunction(isplit->ts,r,rhsfunc,ctx);
150:   VecDestroy(&ralloc);
151:   return(0);
152: }

154: /*@C
155:    TSRHSSplitGetSubTS - Get the sub-TS by split name.

157:    Logically Collective on TS

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

163:    Level: advanced

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

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

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

185:    Logically Collective on TS

187:    Output Parameters:
188: +  n - the number of splits
189: -  subksp - the array of TS contexts

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

195:    Level: advanced

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

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