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: .keywords: TS, TSRHSSplit
 34: @*/
 35: PetscErrorCode TSRHSSplitSetIS(TS ts,const char splitname[],IS is)
 36: {
 37:   TS_RHSSplitLink newsplit,next = ts->tsrhssplit;
 38:   char            prefix[128];
 39:   PetscErrorCode  ierr;


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

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

 71:    Logically Collective on TS

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

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

 80:    Level: intermediate

 82: .seealso: TSRHSSplitSetIS()

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

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

100: /*@C
101:    TSRHSSplitSetRHSFunction - Set the split right-hand-side functions.

103:    Logically Collective on TS

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

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

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

120:  Level: beginner

122: .keywords: TS, timestep, set, ODE, Hamiltonian, Function
123: @*/
124: PetscErrorCode TSRHSSplitSetRHSFunction(TS ts,const char splitname[],Vec r,TSRHSFunction rhsfunc,void *ctx)
125: {
126:   TS_RHSSplitLink isplit;
127:   DM              dmc;
128:   Vec             subvec,ralloc = NULL;
129:   PetscErrorCode  ierr;


135:   /* look up the split */
136:   TSRHSSplitGetRHSSplit(ts,splitname,&isplit);
137:   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);

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

157: /*@C
158:    TSRHSSplitGetSubTS - Get the sub-TS by split name.

160:    Logically Collective on TS

162:    Output Parameters:
163: +  splitname - the number of the split
164: -  subts - the array of TS contexts

166:    Level: advanced

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

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

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

188:    Logically Collective on TS

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

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

198:    Level: advanced

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

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