xref: /petsc/src/ts/impls/explicit/rk/mrk.c (revision e5bffa22434ee5e6fe35b8a1526b6ced76bd6539)
1 /*
2   Code for time stepping with the multi-rate Runge-Kutta method
3 
4   Notes:
5   1) The general system is written as
6      Udot = F(t,U) for the nonsplit version of multi-rate RK,
7      user should give the indexes for both slow and fast components;
8   2) The general system is written as
9      Usdot = Fs(t,Us,Uf)
10      Ufdot = Ff(t,Us,Uf) for multi-rate RK with RHS splits,
11      user should partioned RHS by themselves and also provide the indexes for both slow and fast components.
12 */
13 
14 #include <petsc/private/tsimpl.h>
15 #include <petscdm.h>
16 #include <../src/ts/impls/explicit/rk/rk.h>
17 
18 static TSRKType TSMRKDefault = TSRK2A;
19 
20 static PetscErrorCode TSSetUp_MRKNONSPLIT(TS ts)
21 {
22   TS_RK          *rk = (TS_RK*)ts->data;
23   RKTableau       tab  = rk->tableau;
24   PetscErrorCode ierr;
25 
26   PetscFunctionBegin;
27   ierr = TSRHSSplitGetIS(ts,"slow",&rk->is_slow);CHKERRQ(ierr);
28   ierr = TSRHSSplitGetIS(ts,"fast",&rk->is_fast);CHKERRQ(ierr);
29   if (!rk->is_slow || !rk->is_fast) SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_USER,"Must set up RHSSplits with TSRHSSplitSetIS() using split names 'slow' and 'fast' respectively in order to use multirate RK");
30   ierr = TSRHSSplitGetSubTS(ts,"slow",&rk->subts_slow);CHKERRQ(ierr);
31   ierr = TSRHSSplitGetSubTS(ts,"fast",&rk->subts_fast);CHKERRQ(ierr);
32   if (!rk->subts_slow || !rk->subts_fast) SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_USER,"Must set up the RHSFunctions for 'slow' and 'fast' components using TSRHSSplitSetRHSFunction() or calling TSSetRHSFunction() for each sub-TS");
33   ierr = VecDuplicate(ts->vec_sol,&rk->X0);CHKERRQ(ierr);
34   ierr = VecDuplicateVecs(ts->vec_sol,tab->s,&rk->YdotRHS_slow);CHKERRQ(ierr);
35   PetscFunctionReturn(0);
36 }
37 
38 static PetscErrorCode TSReset_MRKNONSPLIT(TS ts)
39 {
40   TS_RK          *rk = (TS_RK*)ts->data;
41   RKTableau      tab  = rk->tableau;
42   PetscErrorCode ierr;
43 
44   PetscFunctionBegin;
45   ierr = VecDestroy(&rk->X0);CHKERRQ(ierr);
46   ierr = VecDestroyVecs(tab->s,&rk->YdotRHS_slow);CHKERRQ(ierr);
47   PetscFunctionReturn(0);
48 }
49 
50 static PetscErrorCode TSInterpolate_MRKNONSPLIT(TS ts,PetscReal itime,Vec X)
51 {
52   TS_RK            *rk = (TS_RK*)ts->data;
53   PetscInt         s  = rk->tableau->s,p = rk->tableau->p,i,j;
54   PetscReal        h;
55   PetscReal        tt,t;
56   PetscScalar      *b;
57   const PetscReal  *B = rk->tableau->binterp;
58   PetscErrorCode   ierr;
59 
60   PetscFunctionBegin;
61   if (!B) SETERRQ1(PetscObjectComm((PetscObject)ts),PETSC_ERR_SUP,"TSRK %s does not have an interpolation formula",rk->tableau->name);
62 
63   switch (rk->status) {
64     case TS_STEP_INCOMPLETE:
65     case TS_STEP_PENDING:
66       h = ts->time_step;
67       t = (itime - ts->ptime)/h;
68       break;
69     case TS_STEP_COMPLETE:
70       h = ts->ptime - ts->ptime_prev;
71       t = (itime - ts->ptime)/h + 1; /* In the interval [0,1] */
72       break;
73     default: SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_PLIB,"Invalid TSStepStatus");
74   }
75   ierr = PetscMalloc1(s,&b);CHKERRQ(ierr);
76   for (i=0; i<s; i++) b[i] = 0;
77   for (j=0,tt=t; j<p; j++,tt*=t) {
78     for (i=0; i<s; i++) {
79       b[i]  += h * B[i*p+j] * tt;
80     }
81   }
82   ierr = VecCopy(rk->X0,X);CHKERRQ(ierr);
83   ierr = VecMAXPY(X,s,b,rk->YdotRHS_slow);CHKERRQ(ierr);
84   ierr = PetscFree(b);CHKERRQ(ierr);
85   PetscFunctionReturn(0);
86 }
87 
88 static PetscErrorCode TSStep_MRKNONSPLIT(TS ts)
89 {
90   TS              nextlevelts;
91   TS_RK           *rk = (TS_RK*)ts->data;
92   RKTableau       tab  = rk->tableau;
93   Vec             *Y = rk->Y,*YdotRHS = rk->YdotRHS,*YdotRHS_slow = rk->YdotRHS_slow;
94   Vec             stage_slow,sol_slow;   /* vectors store the slow components */
95   Vec             subvec_slow;           /* sub vector to store the slow components */
96   IS              is_slow = rk->is_slow;
97   const PetscInt  s = tab->s;
98   const PetscReal *A = tab->A,*c = tab->c;
99   PetscScalar     *w = rk->work;
100   PetscInt        i,j,k,dtratio = rk->dtratio;
101   PetscReal       next_time_step = ts->time_step,t = ts->ptime,h = ts->time_step;
102   PetscErrorCode  ierr;
103 
104   PetscFunctionBegin;
105   rk->status = TS_STEP_INCOMPLETE;
106   ierr = VecDuplicate(ts->vec_sol,&stage_slow);CHKERRQ(ierr);
107   ierr = VecDuplicate(ts->vec_sol,&sol_slow);CHKERRQ(ierr);
108   ierr = VecCopy(ts->vec_sol,rk->X0);CHKERRQ(ierr);
109   for (i=0; i<s; i++) {
110     rk->stage_time = t + h*c[i];
111     ierr = TSPreStage(ts,rk->stage_time);CHKERRQ(ierr);
112     ierr = VecCopy(ts->vec_sol,Y[i]);CHKERRQ(ierr);
113     for (j=0; j<i; j++) w[j] = h*A[i*s+j];
114     ierr = VecMAXPY(Y[i],i,w,YdotRHS_slow);CHKERRQ(ierr);
115     ierr = TSPostStage(ts,rk->stage_time,i,Y); CHKERRQ(ierr);
116     /* compute the stage RHS */
117     ierr = TSComputeRHSFunction(ts,t+h*c[i],Y[i],YdotRHS_slow[i]);CHKERRQ(ierr);
118   }
119   /* update the slow components in the solution */
120   rk->YdotRHS = YdotRHS_slow;
121   rk->dtratio = 1;
122   ierr = TSEvaluateStep(ts,tab->order,sol_slow,NULL);CHKERRQ(ierr);
123   rk->dtratio = dtratio;
124   rk->YdotRHS = YdotRHS;
125 
126   ierr = TSRHSSplitGetSubTS(rk->subts_fast,"fast",&nextlevelts);CHKERRQ(ierr);
127   if (nextlevelts) {
128     for (k=0; k<rk->dtratio; k++) {
129       ierr = TSRHSSplitGetIS(nextlevelts,"slow",&rk->is_slow);CHKERRQ(ierr);
130       ts->time_step = h/rk->dtratio;
131       ierr = TSStep_MRKNONSPLIT(ts);CHKERRQ(ierr);
132       ts->time_step = h;
133     }
134   } else {
135     for (k=0; k<rk->dtratio; k++) {
136       for (i=0; i<s; i++) {
137         /* update the fast components in the stage value, the slow components will be overwritten, so it is ok to have garbage in the slow components */
138         ierr = VecCopy(ts->vec_sol,Y[i]);CHKERRQ(ierr);
139         for (j=0; j<i; j++) w[j] = h/rk->dtratio*A[i*s+j];
140         ierr = VecMAXPY(Y[i],i,w,YdotRHS);CHKERRQ(ierr);
141         ierr = TSInterpolate_MRKNONSPLIT(ts,t+k*h/rk->dtratio+h/rk->dtratio*c[i],stage_slow);CHKERRQ(ierr);
142         /* update the slow components in the stage value */
143         ierr = VecGetSubVector(stage_slow,rk->is_slow,&subvec_slow);CHKERRQ(ierr);
144         ierr = VecISCopy(Y[i],rk->is_slow,SCATTER_FORWARD,subvec_slow);CHKERRQ(ierr);
145         ierr = VecRestoreSubVector(stage_slow,rk->is_slow,&subvec_slow);CHKERRQ(ierr);
146         /* compute the stage RHS */
147         ierr = TSComputeRHSFunction(ts,t+k*h/rk->dtratio+h/rk->dtratio*c[i],Y[i],YdotRHS[i]);CHKERRQ(ierr);
148       }
149       ierr = TSEvaluateStep(ts,tab->order,ts->vec_sol,NULL);CHKERRQ(ierr);
150     }
151   }
152   /* update the slow components in the solution */
153   ierr = VecGetSubVector(sol_slow,is_slow,&subvec_slow);CHKERRQ(ierr);
154   ierr = VecISCopy(ts->vec_sol,is_slow,SCATTER_FORWARD,subvec_slow);CHKERRQ(ierr);
155   ierr = VecRestoreSubVector(sol_slow,is_slow,&subvec_slow);CHKERRQ(ierr);
156 
157   ts->ptime = t + ts->time_step;
158   ts->time_step = next_time_step;
159   rk->status = TS_STEP_COMPLETE;
160   /* free memory of work vectors */
161   ierr = VecDestroy(&stage_slow);CHKERRQ(ierr);
162   ierr = VecDestroy(&sol_slow);CHKERRQ(ierr);
163   PetscFunctionReturn(0);
164 }
165 
166 static PetscErrorCode TSSetUp_MRKSPLIT(TS ts)
167 {
168   TS_RK          *rk = (TS_RK*)ts->data,*nextlevelrk,*currentlevelrk;
169   TS             nextlevelts;
170   DM             dm,subdm,newdm;
171   PetscErrorCode ierr;
172 
173   PetscFunctionBegin;
174   ierr = TSRHSSplitGetIS(ts,"slow",&rk->is_slow);CHKERRQ(ierr);
175   ierr = TSRHSSplitGetIS(ts,"fast",&rk->is_fast);CHKERRQ(ierr);
176   if (!rk->is_slow || !rk->is_fast) SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_USER,"Must set up RHSSplits with TSRHSSplitSetIS() using split names 'slow' and 'fast' respectively in order to use -ts_type bsi");
177   ierr = TSRHSSplitGetSubTS(ts,"slow",&rk->subts_slow);CHKERRQ(ierr);
178   ierr = TSRHSSplitGetSubTS(ts,"fast",&rk->subts_fast);CHKERRQ(ierr);
179   if (!rk->subts_slow || !rk->subts_fast) SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_USER,"Must set up the RHSFunctions for 'slow' and 'fast' components using TSRHSSplitSetRHSFunction() or calling TSSetRHSFunction() for each sub-TS");
180 
181   /* Only copy */
182   ierr = TSGetDM(ts,&dm);CHKERRQ(ierr);
183   ierr = DMClone(dm,&newdm);CHKERRQ(ierr);
184   ierr = TSGetDM(rk->subts_fast,&subdm);CHKERRQ(ierr);
185   ierr = DMCopyDMTS(subdm,newdm);CHKERRQ(ierr);
186   ierr = DMCopyDMSNES(subdm,newdm);CHKERRQ(ierr);
187   ierr = TSSetDM(rk->subts_fast,newdm);CHKERRQ(ierr);
188   ierr = DMDestroy(&newdm);CHKERRQ(ierr);
189   ierr = DMClone(dm,&newdm);CHKERRQ(ierr);
190   ierr = TSGetDM(rk->subts_slow,&subdm);CHKERRQ(ierr);
191   ierr = DMCopyDMTS(subdm,newdm);CHKERRQ(ierr);
192   ierr = DMCopyDMSNES(subdm,newdm);CHKERRQ(ierr);
193   ierr = TSSetDM(rk->subts_slow,newdm);CHKERRQ(ierr);
194   ierr = DMDestroy(&newdm);CHKERRQ(ierr);
195 
196   ierr = VecDuplicate(ts->vec_sol,&rk->X0);CHKERRQ(ierr);
197   /* The TS at each level share the same tableau, work array, solution vector, stage values and stage derivatives */
198   currentlevelrk = rk;
199   while (currentlevelrk->subts_fast) {
200     /* set up the ts for the slow part */
201     nextlevelts = currentlevelrk->subts_slow;
202     ierr = PetscNewLog(nextlevelts,&nextlevelrk);CHKERRQ(ierr);
203     nextlevelrk->tableau = rk->tableau;
204     nextlevelrk->work = rk->work;
205     nextlevelrk->Y = rk->Y;
206     nextlevelrk->X0 = rk->X0;
207     ierr = PetscMalloc1(rk->tableau->s,&nextlevelrk->YdotRHS);CHKERRQ(ierr);
208     nextlevelts->data = (void*)nextlevelrk;
209     ierr = TSSetSolution(nextlevelts,ts->vec_sol);CHKERRQ(ierr);
210 
211     /* set up the ts for the fast part */
212     nextlevelts = currentlevelrk->subts_fast;
213     ierr = PetscNewLog(nextlevelts,&nextlevelrk);CHKERRQ(ierr);
214     nextlevelrk->tableau = rk->tableau;
215     nextlevelrk->work = rk->work;
216     nextlevelrk->Y = rk->Y;
217     nextlevelrk->X0 = rk->X0;
218     nextlevelrk->dtratio = rk->dtratio;
219     ierr = PetscMalloc1(rk->tableau->s,&nextlevelrk->YdotRHS);CHKERRQ(ierr);
220     ierr = TSRHSSplitGetIS(nextlevelts,"slow",&nextlevelrk->is_slow);CHKERRQ(ierr);
221     ierr = TSRHSSplitGetSubTS(nextlevelts,"slow",&nextlevelrk->subts_slow);CHKERRQ(ierr);
222     ierr = TSRHSSplitGetIS(nextlevelts,"fast",&nextlevelrk->is_fast);CHKERRQ(ierr);
223     ierr = TSRHSSplitGetSubTS(nextlevelts,"fast",&nextlevelrk->subts_fast);CHKERRQ(ierr);
224     nextlevelts->data = (void*)nextlevelrk;
225     ierr = TSSetSolution(nextlevelts,ts->vec_sol);CHKERRQ(ierr);
226 
227     currentlevelrk = nextlevelrk;
228   }
229   PetscFunctionReturn(0);
230 }
231 
232 static PetscErrorCode TSReset_MRKSPLIT(TS ts)
233 {
234   TS_RK          *rk = (TS_RK*)ts->data;
235   PetscErrorCode ierr;
236 
237   PetscFunctionBegin;
238   if (rk->subts_slow) {
239     ierr = TSReset_MRKSPLIT(rk->subts_slow);CHKERRQ(ierr);
240     ierr = PetscFree(rk->subts_slow->data);CHKERRQ(ierr);
241     rk->subts_slow = NULL;
242   }
243   if (rk->subts_fast) {
244     ierr = VecDestroy(&rk->X0);CHKERRQ(ierr);
245     ierr = TSReset_MRKSPLIT(rk->subts_fast);
246     ierr = PetscFree(rk->subts_fast->data);CHKERRQ(ierr);
247     rk->subts_fast = NULL;
248   }
249   ierr = PetscFree(rk->YdotRHS_fast);CHKERRQ(ierr);
250   ierr = PetscFree(rk->YdotRHS_slow);CHKERRQ(ierr);
251   ierr = PetscFree(rk->YdotRHS);CHKERRQ(ierr);
252   PetscFunctionReturn(0);
253 }
254 
255 static PetscErrorCode TSInterpolate_MRKSPLIT(TS ts,PetscReal itime,Vec X)
256 {
257   TS_RK            *rk = (TS_RK*)ts->data;
258   PetscInt         s  = rk->tableau->s,p = rk->tableau->p,i,j;
259   Vec              Yslow;    /* vector holds the slow components in Y[0] */
260   PetscReal        h;
261   PetscReal        tt,t;
262   PetscScalar      *b;
263   const PetscReal  *B = rk->tableau->binterp;
264   PetscErrorCode   ierr;
265 
266   PetscFunctionBegin;
267   if (!B) SETERRQ1(PetscObjectComm((PetscObject)ts),PETSC_ERR_SUP,"TSRK %s does not have an interpolation formula",rk->tableau->name);
268 
269   switch (rk->status) {
270     case TS_STEP_INCOMPLETE:
271     case TS_STEP_PENDING:
272       h = ts->time_step;
273       t = (itime - ts->ptime)/h;
274       break;
275     case TS_STEP_COMPLETE:
276       h = ts->ptime - ts->ptime_prev;
277       t = (itime - ts->ptime)/h + 1; /* In the interval [0,1] */
278       break;
279     default: SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_PLIB,"Invalid TSStepStatus");
280   }
281   ierr = PetscMalloc1(s,&b);CHKERRQ(ierr);
282   for (i=0; i<s; i++) b[i] = 0;
283   for (j=0,tt=t; j<p; j++,tt*=t) {
284     for (i=0; i<s; i++) {
285       b[i]  += h * B[i*p+j] * tt;
286     }
287   }
288   ierr = VecGetSubVector(rk->X0,rk->is_slow,&Yslow);CHKERRQ(ierr);
289   ierr = VecCopy(Yslow,X);CHKERRQ(ierr);
290   ierr = VecMAXPY(X,s,b,((TS_RK*)rk->subts_slow->data)->YdotRHS);CHKERRQ(ierr);
291   ierr = VecRestoreSubVector(rk->X0,rk->is_slow,&Yslow);CHKERRQ(ierr);
292   ierr = PetscFree(b);CHKERRQ(ierr);
293   PetscFunctionReturn(0);
294 }
295 
296 /*
297  This is for partitioned RHS multirate RK method
298  The step completion formula is
299 
300  x1 = x0 + h b^T YdotRHS
301 
302 */
303 static PetscErrorCode TSEvaluateStep_MRKSPLIT(TS ts,PetscInt order,Vec X,PetscBool *done)
304 {
305   TS_RK           *rk = (TS_RK*)ts->data;
306   RKTableau       tab  = rk->tableau;
307   Vec             Xslow,Xfast;                  /* subvectors of X which store slow components and fast components respectively */
308   PetscScalar     *w = rk->work;
309   PetscReal       h = ts->time_step;
310   PetscInt        s = tab->s,j;
311   PetscErrorCode  ierr;
312 
313   PetscFunctionBegin;
314   ierr = VecCopy(ts->vec_sol,X);CHKERRQ(ierr);
315   if (rk->slow) {
316     for (j=0; j<s; j++) w[j] = h*tab->b[j];
317     ierr = VecGetSubVector(ts->vec_sol,rk->is_slow,&Xslow);CHKERRQ(ierr);
318     ierr = VecMAXPY(Xslow,s,w,((TS_RK*)rk->subts_slow->data)->YdotRHS);CHKERRQ(ierr);
319     ierr = VecRestoreSubVector(ts->vec_sol,rk->is_slow,&Xslow);CHKERRQ(ierr);;
320   } else {
321     for (j=0; j<s; j++) w[j] = h/rk->dtratio*tab->b[j];
322     ierr = VecGetSubVector(X,rk->is_fast,&Xfast);CHKERRQ(ierr);
323     ierr = VecMAXPY(Xfast,s,w,((TS_RK*)rk->subts_fast->data)->YdotRHS);CHKERRQ(ierr);
324     ierr = VecRestoreSubVector(X,rk->is_fast,&Xfast);CHKERRQ(ierr);
325   }
326   PetscFunctionReturn(0);
327 }
328 
329 static PetscErrorCode TSStep_MRKSPLIT(TS ts)
330 {
331   TS_RK             *rk = (TS_RK*)ts->data;
332   RKTableau         tab = rk->tableau;
333   Vec               *Y = rk->Y,*YdotRHS = rk->YdotRHS;
334   Vec               *YdotRHS_fast,*YdotRHS_slow;
335   Vec               Yslow,Yfast;                         /* subvectors store the stges of slow components and fast components respectively                           */
336   const PetscInt    s = tab->s;
337   const PetscReal   *A = tab->A,*c = tab->c;
338   PetscScalar       *w = rk->work;
339   PetscInt          i,j,k;
340   PetscReal         next_time_step = ts->time_step,t = ts->ptime,h = ts->time_step;
341   PetscErrorCode    ierr;
342 
343   PetscFunctionBegin;
344   rk->status = TS_STEP_INCOMPLETE;
345   YdotRHS_fast = ((TS_RK*)rk->subts_fast->data)->YdotRHS;
346   YdotRHS_slow = ((TS_RK*)rk->subts_slow->data)->YdotRHS;
347   for (i=0; i<s; i++) {
348     ierr = VecGetSubVector(YdotRHS[i],rk->is_slow,&YdotRHS_slow[i]);CHKERRQ(ierr);
349     ierr = VecGetSubVector(YdotRHS[i],rk->is_fast,&YdotRHS_fast[i]);CHKERRQ(ierr);
350   }
351   ierr = VecCopy(ts->vec_sol,rk->X0);CHKERRQ(ierr);
352   /* propogate both slow and fast components using large time steps */
353   for (i=0; i<s; i++) {
354     rk->stage_time = t + h*c[i];
355     ierr = TSPreStage(ts,rk->stage_time);CHKERRQ(ierr);
356     ierr = VecCopy(ts->vec_sol,Y[i]);CHKERRQ(ierr);
357     /*ierr = VecCopy(ts->vec_sol,Yc[i]);CHKERRQ(ierr);*/
358     ierr = VecGetSubVector(Y[i],rk->is_fast,&Yfast);CHKERRQ(ierr);
359     ierr = VecGetSubVector(Y[i],rk->is_slow,&Yslow);CHKERRQ(ierr);
360     for (j=0; j<i; j++) w[j] = h*A[i*s+j];
361     ierr = VecMAXPY(Yfast,i,w,YdotRHS_fast);CHKERRQ(ierr);
362     ierr = VecMAXPY(Yslow,i,w,YdotRHS_slow);CHKERRQ(ierr);
363     ierr = VecRestoreSubVector(Y[i],rk->is_fast,&Yfast);CHKERRQ(ierr);
364     ierr = VecRestoreSubVector(Y[i],rk->is_slow,&Yslow);CHKERRQ(ierr);
365     ierr = TSPostStage(ts,rk->stage_time,i,Y); CHKERRQ(ierr);
366     ierr = TSComputeRHSFunction(rk->subts_slow,t+h*c[i],Y[i],YdotRHS_slow[i]);CHKERRQ(ierr);
367     ierr = TSComputeRHSFunction(rk->subts_fast,t+h*c[i],Y[i],YdotRHS_fast[i]);CHKERRQ(ierr);
368   }
369   rk->slow = PETSC_TRUE;
370   /* update the value of slow components when using slow time step */
371   ierr = TSEvaluateStep_MRKSPLIT(ts,tab->order,ts->vec_sol,NULL);CHKERRQ(ierr);
372 
373   if (((TS_RK*)rk->subts_fast->data)->subts_fast) {
374     rk->subts_fast->ptime = t;
375     rk->subts_fast->time_step = h/rk->dtratio;
376     for (k=0; k<rk->dtratio; k++) {
377       ierr = TSStep_MRKSPLIT(rk->subts_fast);CHKERRQ(ierr);
378     }
379   } else {
380     for (k=0; k<rk->dtratio; k++) {
381       /* propogate fast component using small time steps */
382       for (i=0; i<s; i++) {
383         rk->stage_time = t + h/rk->dtratio*c[i];
384         ierr = TSPreStage(ts,rk->stage_time);CHKERRQ(ierr);
385         ierr = VecCopy(ts->vec_sol,Y[i]);CHKERRQ(ierr);
386         /* stage value for fast components */
387         for (j=0; j<i; j++) w[j] = h/rk->dtratio*A[i*s+j];
388         ierr = VecGetSubVector(Y[i],rk->is_fast,&Yfast);CHKERRQ(ierr);
389         ierr = VecMAXPY(Yfast,i,w,YdotRHS_fast);CHKERRQ(ierr);
390         ierr = VecRestoreSubVector(Y[i],rk->is_fast,&Yfast);CHKERRQ(ierr);
391         /* stage value for slow components */
392         ierr = VecGetSubVector(Y[i],rk->is_slow,&Yslow);CHKERRQ(ierr);
393         ierr = TSInterpolate_MRKSPLIT(ts,t+k*h/rk->dtratio+h/rk->dtratio*c[i],Yslow);CHKERRQ(ierr);
394         ierr = VecRestoreSubVector(Y[i],rk->is_slow,&Yslow);CHKERRQ(ierr);
395         ierr = TSPostStage(ts,rk->stage_time,i,Y);CHKERRQ(ierr);
396         /* compute the stage RHS for fast components */
397         ierr = TSComputeRHSFunction(rk->subts_fast,t+k*h*rk->dtratio+h/rk->dtratio*c[i],Y[i],YdotRHS_fast[i]);CHKERRQ(ierr);
398       }
399       /* update the value of fast components when using fast time step */
400       rk->slow = PETSC_FALSE;
401       ierr = TSEvaluateStep_MRKSPLIT(ts,tab->order,ts->vec_sol,NULL);CHKERRQ(ierr);
402     }
403   }
404   for (i=0; i<s; i++) {
405     ierr = VecRestoreSubVector(YdotRHS[i],rk->is_slow,&YdotRHS_slow[i]);CHKERRQ(ierr);
406     ierr = VecRestoreSubVector(YdotRHS[i],rk->is_fast,&YdotRHS_fast[i]);CHKERRQ(ierr);
407   }
408   ts->ptime = t + ts->time_step;
409   ts->time_step = next_time_step;
410   rk->status = TS_STEP_COMPLETE;
411   PetscFunctionReturn(0);
412 }
413 
414 /*@C
415   TSRKSetMultirateType - Set the type of RK Multirate scheme
416 
417   Logically collective
418 
419   Input Parameter:
420 +  ts - timestepping context
421 -  mrktype - type of MRK-scheme
422 
423   Options Database:
424 .   -ts_rk_multiarte_type - <none,nonsplit,split>
425 
426   Level: intermediate
427 @*/
428 PetscErrorCode TSRKSetMultirateType(TS ts, TSMRKType mrktype)
429 {
430   TS_RK          *rk = (TS_RK*)ts->data;
431   PetscErrorCode ierr;
432 
433   PetscFunctionBegin;
434   PetscValidHeaderSpecific(ts,TS_CLASSID,1);
435   switch(mrktype){
436     case TSMRKNONE:
437       break;
438     case TSMRKNONSPLIT:
439       ts->ops->step           = TSStep_MRKNONSPLIT;
440       ts->ops->interpolate    = TSInterpolate_MRKNONSPLIT;
441       rk->dtratio             = 2;
442       ierr = TSRKSetType(ts,TSMRKDefault);CHKERRQ(ierr);
443       ierr = PetscObjectComposeFunction((PetscObject)ts,"TSSetUp_MRKNONSPLIT_C",TSSetUp_MRKNONSPLIT);CHKERRQ(ierr);
444       ierr = PetscObjectComposeFunction((PetscObject)ts,"TSReset_MRKNONSPLIT_C",TSReset_MRKNONSPLIT);CHKERRQ(ierr);
445       break;
446     case TSMRKSPLIT:
447       ts->ops->step           = TSStep_MRKSPLIT;
448       ts->ops->evaluatestep   = TSEvaluateStep_MRKSPLIT;
449       ts->ops->interpolate    = TSInterpolate_MRKSPLIT;
450       rk->dtratio             = 2;
451       ierr = TSRKSetType(ts,TSMRKDefault);CHKERRQ(ierr);
452       ierr = PetscObjectComposeFunction((PetscObject)ts,"TSSetUp_MRKSPLIT_C",TSSetUp_MRKSPLIT);CHKERRQ(ierr);
453       ierr = PetscObjectComposeFunction((PetscObject)ts,"TSReset_MRKSPLIT_C",TSReset_MRKSPLIT);CHKERRQ(ierr);
454       break;
455     default :
456       SETERRQ1(PetscObjectComm((PetscObject)ts),PETSC_ERR_ARG_UNKNOWN_TYPE,"Unknown type '%s'",mrktype);
457   }
458   PetscFunctionReturn(0);
459 }
460 
461