xref: /petsc/src/ts/impls/explicit/rk/mrk.c (revision 19c2959aedcbb17969a32880b9fa35aa707cd651)
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   rk->subts_current = rk->subts_fast;
36   PetscFunctionReturn(0);
37 }
38 
39 static PetscErrorCode TSReset_MRKNONSPLIT(TS ts)
40 {
41   TS_RK          *rk = (TS_RK*)ts->data;
42   RKTableau      tab = rk->tableau;
43   PetscErrorCode ierr;
44 
45   PetscFunctionBegin;
46   ierr = VecDestroy(&rk->X0);CHKERRQ(ierr);
47   ierr = VecDestroyVecs(tab->s,&rk->YdotRHS_slow);CHKERRQ(ierr);
48   PetscFunctionReturn(0);
49 }
50 
51 static PetscErrorCode TSInterpolate_MRKNONSPLIT(TS ts,PetscReal itime,Vec X)
52 {
53   TS_RK            *rk = (TS_RK*)ts->data;
54   PetscInt         s = rk->tableau->s,p = rk->tableau->p,i,j;
55   PetscReal        h = ts->time_step;
56   PetscReal        tt,t;
57   PetscScalar      *b;
58   const PetscReal  *B = rk->tableau->binterp;
59   PetscErrorCode   ierr;
60 
61   PetscFunctionBegin;
62   if (!B) SETERRQ1(PetscObjectComm((PetscObject)ts),PETSC_ERR_SUP,"TSRK %s does not have an interpolation formula",rk->tableau->name);
63   t = (itime - rk->ptime)/h;
64   ierr = PetscMalloc1(s,&b);CHKERRQ(ierr);
65   for (i=0; i<s; i++) b[i] = 0;
66   for (j=0,tt=t; j<p; j++,tt*=t) {
67     for (i=0; i<s; i++) {
68       b[i]  += h * B[i*p+j] * tt;
69     }
70   }
71   ierr = VecCopy(rk->X0,X);CHKERRQ(ierr);
72   ierr = VecMAXPY(X,s,b,rk->YdotRHS_slow);CHKERRQ(ierr);
73   ierr = PetscFree(b);CHKERRQ(ierr);
74   PetscFunctionReturn(0);
75 }
76 
77 static PetscErrorCode TSStepRefine_MRKNONSPLIT(TS ts)
78 {
79   TS              previousts,subts;
80   TS_RK           *rk = (TS_RK*)ts->data;
81   RKTableau       tab = rk->tableau;
82   Vec             *Y = rk->Y,*YdotRHS = rk->YdotRHS;
83   Vec             vec_fast,subvec_fast;
84   const PetscInt  s = tab->s;
85   const PetscReal *A = tab->A,*c = tab->c;
86   PetscScalar     *w = rk->work;
87   PetscInt        i,j,k;
88   PetscReal       t = ts->ptime,h = ts->time_step;
89   PetscErrorCode  ierr;
90 
91   ierr = VecDuplicate(ts->vec_sol,&vec_fast);CHKERRQ(ierr);
92   previousts = rk->subts_current;
93   ierr = TSRHSSplitGetSubTS(rk->subts_current,"fast",&subts);CHKERRQ(ierr);
94   ierr = TSRHSSplitGetSubTS(subts,"fast",&subts);CHKERRQ(ierr);
95   for (k=0; k<rk->dtratio; k++) {
96     for (i=0; i<s; i++) {
97       ierr = TSInterpolate_MRKNONSPLIT(ts,t+k*h/rk->dtratio+h/rk->dtratio*c[i],Y[i]);CHKERRQ(ierr);
98       for (j=0; j<i; j++) w[j] = h/rk->dtratio*A[i*s+j];
99       /* update the fast components in the stage value, the slow components will be ignored, so we do not care the slow part in vec_fast */
100       ierr = VecCopy(ts->vec_sol,vec_fast);CHKERRQ(ierr);
101       ierr = VecMAXPY(vec_fast,i,w,YdotRHS);CHKERRQ(ierr);
102       /* update the fast components in the stage value */
103       ierr = VecGetSubVector(vec_fast,rk->is_fast,&subvec_fast);CHKERRQ(ierr);
104       ierr = VecISCopy(Y[i],rk->is_fast,SCATTER_FORWARD,subvec_fast);CHKERRQ(ierr);
105       ierr = VecRestoreSubVector(vec_fast,rk->is_fast,&subvec_fast);CHKERRQ(ierr);
106       /* compute the stage RHS */
107       ierr = TSComputeRHSFunction(ts,t+k*h/rk->dtratio+h/rk->dtratio*c[i],Y[i],YdotRHS[i]);CHKERRQ(ierr);
108     }
109     ierr = VecCopy(ts->vec_sol,vec_fast);CHKERRQ(ierr);
110     ierr = TSEvaluateStep(ts,tab->order,vec_fast,NULL);CHKERRQ(ierr);
111     /* update the fast components in the solution */
112     ierr = VecGetSubVector(vec_fast,rk->is_fast,&subvec_fast);CHKERRQ(ierr);
113     ierr = VecISCopy(ts->vec_sol,rk->is_fast,SCATTER_FORWARD,subvec_fast);CHKERRQ(ierr);
114     ierr = VecRestoreSubVector(vec_fast,rk->is_fast,&subvec_fast);CHKERRQ(ierr);
115 
116     if (subts) {
117       Vec *YdotRHS_copy;
118       ierr = VecDuplicateVecs(ts->vec_sol,s,&YdotRHS_copy);CHKERRQ(ierr);
119       rk->subts_current = rk->subts_fast;
120       ts->ptime = t+k*h/rk->dtratio;
121       ts->time_step = h/rk->dtratio;
122       ierr = TSRHSSplitGetIS(rk->subts_current,"fast",&rk->is_fast);CHKERRQ(ierr);
123       for (i=0; i<s; i++) {
124         ierr = VecCopy(rk->YdotRHS_slow[i],YdotRHS_copy[i]);CHKERRQ(ierr);
125         ierr = VecCopy(YdotRHS[i],rk->YdotRHS_slow[i]);CHKERRQ(ierr);
126       }
127 
128       ierr = TSStepRefine_MRKNONSPLIT(ts);CHKERRQ(ierr);
129 
130       rk->subts_current = previousts;
131       ts->ptime = t;
132       ts->time_step = h;
133       ierr = TSRHSSplitGetIS(previousts,"fast",&rk->is_fast);CHKERRQ(ierr);
134       for (i=0; i<s; i++) {
135         ierr = VecCopy(YdotRHS_copy[i],rk->YdotRHS_slow[i]);CHKERRQ(ierr);
136       }
137       ierr = VecDestroyVecs(s,&YdotRHS_copy);CHKERRQ(ierr);
138     }
139   }
140   ierr = VecDestroy(&vec_fast);CHKERRQ(ierr);
141   PetscFunctionReturn(0);
142 }
143 
144 static PetscErrorCode TSStep_MRKNONSPLIT(TS ts)
145 {
146   TS_RK           *rk = (TS_RK*)ts->data;
147   RKTableau       tab = rk->tableau;
148   Vec             *Y = rk->Y,*YdotRHS = rk->YdotRHS,*YdotRHS_slow = rk->YdotRHS_slow;
149   Vec             stage_slow,sol_slow; /* vectors store the slow components */
150   Vec             subvec_slow; /* sub vector to store the slow components */
151   IS              is_slow = rk->is_slow;
152   const PetscInt  s = tab->s;
153   const PetscReal *A = tab->A,*c = tab->c;
154   PetscScalar     *w = rk->work;
155   PetscInt        i,j,dtratio = rk->dtratio;
156   PetscReal       next_time_step = ts->time_step,t = ts->ptime,h = ts->time_step;
157   PetscErrorCode  ierr;
158 
159   PetscFunctionBegin;
160   rk->status = TS_STEP_INCOMPLETE;
161   ierr = VecDuplicate(ts->vec_sol,&stage_slow);CHKERRQ(ierr);
162   ierr = VecDuplicate(ts->vec_sol,&sol_slow);CHKERRQ(ierr);
163   ierr = VecCopy(ts->vec_sol,rk->X0);CHKERRQ(ierr);
164   for (i=0; i<s; i++) {
165     rk->stage_time = t + h*c[i];
166     ierr = TSPreStage(ts,rk->stage_time);CHKERRQ(ierr);
167     ierr = VecCopy(ts->vec_sol,Y[i]);CHKERRQ(ierr);
168     for (j=0; j<i; j++) w[j] = h*A[i*s+j];
169     ierr = VecMAXPY(Y[i],i,w,YdotRHS_slow);CHKERRQ(ierr);
170     ierr = TSPostStage(ts,rk->stage_time,i,Y); CHKERRQ(ierr);
171     /* compute the stage RHS */
172     ierr = TSComputeRHSFunction(ts,t+h*c[i],Y[i],YdotRHS_slow[i]);CHKERRQ(ierr);
173   }
174   /* update the slow components in the solution */
175   rk->YdotRHS = YdotRHS_slow;
176   rk->dtratio = 1;
177   ierr = TSEvaluateStep(ts,tab->order,sol_slow,NULL);CHKERRQ(ierr);
178   rk->dtratio = dtratio;
179   rk->YdotRHS = YdotRHS;
180   /* update the slow components in the solution */
181   ierr = VecGetSubVector(sol_slow,is_slow,&subvec_slow);CHKERRQ(ierr);
182   ierr = VecISCopy(ts->vec_sol,is_slow,SCATTER_FORWARD,subvec_slow);CHKERRQ(ierr);
183   ierr = VecRestoreSubVector(sol_slow,is_slow,&subvec_slow);CHKERRQ(ierr);
184 
185   rk->subts_current = ts;
186   rk->ptime = t;
187   rk->time_step = h;
188   ierr = TSStepRefine_MRKNONSPLIT(ts);CHKERRQ(ierr);
189 
190   ts->ptime = t + ts->time_step;
191   ts->time_step = next_time_step;
192   rk->status = TS_STEP_COMPLETE;
193 
194   /* free memory of work vectors */
195   ierr = VecDestroy(&stage_slow);CHKERRQ(ierr);
196   ierr = VecDestroy(&sol_slow);CHKERRQ(ierr);
197   PetscFunctionReturn(0);
198 }
199 
200 /*
201   Copy DM from tssrc to tsdest, while keeping the original DMTS and DMSNES in tsdest.
202 */
203 static PetscErrorCode TSCopyDM(TS tssrc,TS tsdest)
204 {
205   DM             newdm,dmsrc,dmdest;
206   PetscErrorCode ierr;
207 
208   PetscFunctionBegin;
209   ierr = TSGetDM(tssrc,&dmsrc);CHKERRQ(ierr);
210   ierr = DMClone(dmsrc,&newdm);CHKERRQ(ierr);
211   ierr = TSGetDM(tsdest,&dmdest);CHKERRQ(ierr);
212   ierr = DMCopyDMTS(dmdest,newdm);CHKERRQ(ierr);
213   ierr = DMCopyDMSNES(dmdest,newdm);CHKERRQ(ierr);
214   ierr = TSSetDM(tsdest,newdm);CHKERRQ(ierr);
215   ierr = DMDestroy(&newdm);CHKERRQ(ierr);
216   PetscFunctionReturn(0);
217 }
218 
219 static PetscErrorCode TSSetUp_MRKSPLIT(TS ts)
220 {
221   TS_RK          *rk = (TS_RK*)ts->data,*nextlevelrk,*currentlevelrk;
222   TS             nextlevelts;
223   Vec            X0;
224   PetscErrorCode ierr;
225 
226   PetscFunctionBegin;
227   ierr = TSRHSSplitGetIS(ts,"slow",&rk->is_slow);CHKERRQ(ierr);
228   ierr = TSRHSSplitGetIS(ts,"fast",&rk->is_fast);CHKERRQ(ierr);
229   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");
230   ierr = TSRHSSplitGetSubTS(ts,"slow",&rk->subts_slow);CHKERRQ(ierr);
231   ierr = TSRHSSplitGetSubTS(ts,"fast",&rk->subts_fast);CHKERRQ(ierr);
232   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");
233 
234   ierr = VecDuplicate(ts->vec_sol,&X0);CHKERRQ(ierr);
235   /* The TS at each level share the same tableau, work array, solution vector, stage values and stage derivatives */
236   currentlevelrk = rk;
237   while (currentlevelrk->subts_fast) {
238     ierr = PetscMalloc1(rk->tableau->s,&currentlevelrk->YdotRHS_fast);CHKERRQ(ierr);
239     ierr = PetscMalloc1(rk->tableau->s,&currentlevelrk->YdotRHS_slow);CHKERRQ(ierr);
240     ierr = PetscObjectReference((PetscObject)X0);CHKERRQ(ierr);
241     currentlevelrk->X0 = X0;
242     currentlevelrk->ts_root = ts;
243 
244     /* set up the ts for the slow part */
245     nextlevelts = currentlevelrk->subts_slow;
246     ierr = PetscNewLog(nextlevelts,&nextlevelrk);CHKERRQ(ierr);
247     nextlevelrk->tableau = rk->tableau;
248     nextlevelrk->work = rk->work;
249     nextlevelrk->Y = rk->Y;
250     nextlevelrk->YdotRHS = rk->YdotRHS;
251     nextlevelts->data = (void*)nextlevelrk;
252     ierr = TSCopyDM(ts,nextlevelts);CHKERRQ(ierr);
253     ierr = TSSetSolution(nextlevelts,ts->vec_sol);CHKERRQ(ierr);
254 
255     /* set up the ts for the fast part */
256     nextlevelts = currentlevelrk->subts_fast;
257     ierr = PetscNewLog(nextlevelts,&nextlevelrk);CHKERRQ(ierr);
258     nextlevelrk->tableau = rk->tableau;
259     nextlevelrk->work = rk->work;
260     nextlevelrk->Y = rk->Y;
261     nextlevelrk->YdotRHS = rk->YdotRHS;
262     nextlevelrk->dtratio = rk->dtratio;
263     ierr = TSRHSSplitGetIS(nextlevelts,"slow",&nextlevelrk->is_slow);CHKERRQ(ierr);
264     ierr = TSRHSSplitGetSubTS(nextlevelts,"slow",&nextlevelrk->subts_slow);CHKERRQ(ierr);
265     ierr = TSRHSSplitGetIS(nextlevelts,"fast",&nextlevelrk->is_fast);CHKERRQ(ierr);
266     ierr = TSRHSSplitGetSubTS(nextlevelts,"fast",&nextlevelrk->subts_fast);CHKERRQ(ierr);
267     nextlevelts->data = (void*)nextlevelrk;
268     ierr = TSCopyDM(ts,nextlevelts);CHKERRQ(ierr);
269     ierr = TSSetSolution(nextlevelts,ts->vec_sol);CHKERRQ(ierr);
270 
271     currentlevelrk = nextlevelrk;
272   }
273   ierr = VecDestroy(&X0);CHKERRQ(ierr);
274   PetscFunctionReturn(0);
275 }
276 
277 static PetscErrorCode TSReset_MRKSPLIT(TS ts)
278 {
279   TS_RK          *rk = (TS_RK*)ts->data;
280   PetscErrorCode ierr;
281 
282   PetscFunctionBegin;
283   if (rk->subts_slow) {
284     ierr = PetscFree(rk->subts_slow->data);CHKERRQ(ierr);
285     rk->subts_slow = NULL;
286   }
287   if (rk->subts_fast) {
288     ierr = PetscFree(rk->YdotRHS_fast);CHKERRQ(ierr);
289     ierr = PetscFree(rk->YdotRHS_slow);CHKERRQ(ierr);
290     ierr = VecDestroy(&rk->X0);CHKERRQ(ierr);
291     ierr = TSReset_MRKSPLIT(rk->subts_fast);
292     ierr = PetscFree(rk->subts_fast->data);CHKERRQ(ierr);
293     rk->subts_fast = NULL;
294   }
295   PetscFunctionReturn(0);
296 }
297 
298 static PetscErrorCode TSInterpolate_MRKSPLIT(TS ts,PetscReal itime,Vec X)
299 {
300   TS_RK           *rk = (TS_RK*)ts->data;
301   Vec             Xslow;
302   PetscInt        s = rk->tableau->s,p = rk->tableau->p,i,j;
303   PetscReal       h;
304   PetscReal       tt,t;
305   PetscScalar     *b;
306   const PetscReal *B = rk->tableau->binterp;
307   PetscErrorCode  ierr;
308 
309   PetscFunctionBegin;
310   if (!B) SETERRQ1(PetscObjectComm((PetscObject)ts),PETSC_ERR_SUP,"TSRK %s does not have an interpolation formula",rk->tableau->name);
311 
312   switch (rk->status) {
313     case TS_STEP_INCOMPLETE:
314     case TS_STEP_PENDING:
315       h = ts->time_step;
316       t = (itime - ts->ptime)/h;
317       break;
318     case TS_STEP_COMPLETE:
319       h = ts->ptime - ts->ptime_prev;
320       t = (itime - ts->ptime)/h + 1; /* In the interval [0,1] */
321       break;
322     default: SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_PLIB,"Invalid TSStepStatus");
323   }
324   ierr = PetscMalloc1(s,&b);CHKERRQ(ierr);
325   for (i=0; i<s; i++) b[i] = 0;
326   for (j=0,tt=t; j<p; j++,tt*=t) {
327     for (i=0; i<s; i++) {
328       b[i]  += h * B[i*p+j] * tt;
329     }
330   }
331   for (i=0; i<s; i++) {
332     ierr = VecGetSubVector(rk->YdotRHS[i],rk->is_slow,&rk->YdotRHS_slow[i]);CHKERRQ(ierr);
333   }
334   ierr = VecGetSubVector(X,rk->is_slow,&Xslow);CHKERRQ(ierr);
335   ierr = VecISCopy(rk->X0,rk->is_slow,SCATTER_REVERSE,Xslow);CHKERRQ(ierr);
336   ierr = VecMAXPY(Xslow,s,b,rk->YdotRHS_slow);CHKERRQ(ierr);
337   ierr = VecRestoreSubVector(X,rk->is_slow,&Xslow);CHKERRQ(ierr);
338   for (i=0; i<s; i++) {
339     ierr = VecRestoreSubVector(rk->YdotRHS[i],rk->is_slow,&rk->YdotRHS_slow[i]);CHKERRQ(ierr);
340   }
341   ierr = PetscFree(b);CHKERRQ(ierr);
342   PetscFunctionReturn(0);
343 }
344 
345 /*
346  This is for partitioned RHS multirate RK method
347  The step completion formula is
348 
349  x1 = x0 + h b^T YdotRHS
350 
351 */
352 static PetscErrorCode TSEvaluateStep_MRKSPLIT(TS ts,PetscInt order,Vec X,PetscBool *done)
353 {
354   TS_RK          *rk = (TS_RK*)ts->data;
355   RKTableau      tab = rk->tableau;
356   Vec            Xslow,Xfast;                  /* subvectors of X which store slow components and fast components respectively */
357   PetscScalar    *w = rk->work;
358   PetscReal      h = ts->time_step;
359   PetscInt       s = tab->s,j;
360   PetscErrorCode ierr;
361 
362   PetscFunctionBegin;
363   ierr = VecCopy(ts->vec_sol,X);CHKERRQ(ierr);
364   if (rk->slow) {
365     for (j=0; j<s; j++) w[j] = h*tab->b[j];
366     ierr = VecGetSubVector(ts->vec_sol,rk->is_slow,&Xslow);CHKERRQ(ierr);
367     ierr = VecMAXPY(Xslow,s,w,rk->YdotRHS_slow);CHKERRQ(ierr);
368     ierr = VecRestoreSubVector(ts->vec_sol,rk->is_slow,&Xslow);CHKERRQ(ierr);;
369   } else {
370     for (j=0; j<s; j++) w[j] = h/rk->dtratio*tab->b[j];
371     ierr = VecGetSubVector(X,rk->is_fast,&Xfast);CHKERRQ(ierr);
372     ierr = VecMAXPY(Xfast,s,w,rk->YdotRHS_fast);CHKERRQ(ierr);
373     ierr = VecRestoreSubVector(X,rk->is_fast,&Xfast);CHKERRQ(ierr);
374   }
375   PetscFunctionReturn(0);
376 }
377 
378 static PetscErrorCode TSStepRefine_MRKSPLIT(TS ts)
379 {
380   TS_RK           *rk = (TS_RK*)ts->data;
381   TS              subts_fast = rk->subts_fast,currentlevelts;
382   TS_RK           *subrk_fast = (TS_RK*)subts_fast->data;
383   RKTableau       tab = rk->tableau;
384   Vec             *Y = rk->Y;
385   Vec             *YdotRHS = rk->YdotRHS,*YdotRHS_fast = rk->YdotRHS_fast;
386   Vec             Yfast,Xfast;
387   const PetscInt  s = tab->s;
388   const PetscReal *A = tab->A,*c = tab->c;
389   PetscScalar     *w = rk->work;
390   PetscInt        i,j,k;
391   PetscReal       t = ts->ptime,h = ts->time_step;
392   PetscErrorCode  ierr;
393 
394   for (k=0; k<rk->dtratio; k++) {
395     ierr = VecGetSubVector(ts->vec_sol,rk->is_fast,&Xfast);CHKERRQ(ierr);
396     for (i=0; i<s; i++) {
397       ierr = VecGetSubVector(YdotRHS[i],rk->is_fast,&YdotRHS_fast[i]);CHKERRQ(ierr);
398     }
399     /* propogate fast component using small time steps */
400     for (i=0; i<s; i++) {
401       /* stage value for slow components */
402       ierr = TSInterpolate_MRKSPLIT(rk->ts_root,t+k*h/rk->dtratio+h/rk->dtratio*c[i],Y[i]);CHKERRQ(ierr);
403       currentlevelts = rk->ts_root;
404       while (currentlevelts != ts) { /* all the slow parts need to be interpolated separated */
405         currentlevelts = ((TS_RK*)currentlevelts->data)->subts_fast;
406         ierr = TSInterpolate_MRKSPLIT(currentlevelts,t+k*h/rk->dtratio+h/rk->dtratio*c[i],Y[i]);CHKERRQ(ierr);
407       }
408       for (j=0; j<i; j++) w[j] = h/rk->dtratio*A[i*s+j];
409       subrk_fast->stage_time = t + h/rk->dtratio*c[i];
410       ierr = TSPreStage(subts_fast,subrk_fast->stage_time);CHKERRQ(ierr);
411       /* stage value for fast components */
412       ierr = VecGetSubVector(Y[i],rk->is_fast,&Yfast);CHKERRQ(ierr);
413       ierr = VecCopy(Xfast,Yfast);CHKERRQ(ierr);
414       ierr = VecMAXPY(Yfast,i,w,YdotRHS_fast);CHKERRQ(ierr);
415       ierr = VecRestoreSubVector(Y[i],rk->is_fast,&Yfast);CHKERRQ(ierr);
416       ierr = TSPostStage(subts_fast,subrk_fast->stage_time,i,Y);CHKERRQ(ierr);
417       /* compute the stage RHS for fast components */
418       ierr = TSComputeRHSFunction(subts_fast,t+k*h*rk->dtratio+h/rk->dtratio*c[i],Y[i],YdotRHS_fast[i]);CHKERRQ(ierr);
419     }
420     ierr = VecRestoreSubVector(ts->vec_sol,rk->is_fast,&Xfast);CHKERRQ(ierr);
421     /* update the value of fast components using fast time step */
422     rk->slow = PETSC_FALSE;
423     ierr = TSEvaluateStep_MRKSPLIT(ts,tab->order,ts->vec_sol,NULL);CHKERRQ(ierr);
424     for (i=0; i<s; i++) {
425       ierr = VecRestoreSubVector(YdotRHS[i],rk->is_fast,&YdotRHS_fast[i]);CHKERRQ(ierr);
426     }
427 
428     if (subrk_fast->subts_fast) {
429       subts_fast->ptime = t+k*h/rk->dtratio;
430       subts_fast->time_step = h/rk->dtratio;
431       ierr = TSStepRefine_MRKSPLIT(subts_fast);CHKERRQ(ierr);
432     }
433     /* update the fast components of the solution */
434     ierr = VecGetSubVector(ts->vec_sol,rk->is_fast,&Xfast);CHKERRQ(ierr);
435     ierr = VecISCopy(rk->X0,rk->is_fast,SCATTER_FORWARD,Xfast);CHKERRQ(ierr);
436     ierr = VecRestoreSubVector(ts->vec_sol,rk->is_fast,&Xfast);CHKERRQ(ierr);
437   }
438   PetscFunctionReturn(0);
439 }
440 
441 static PetscErrorCode TSStep_MRKSPLIT(TS ts)
442 {
443   TS_RK           *rk = (TS_RK*)ts->data;
444   RKTableau       tab = rk->tableau;
445   Vec             *Y = rk->Y,*YdotRHS = rk->YdotRHS;
446   Vec             *YdotRHS_fast = rk->YdotRHS_fast,*YdotRHS_slow = rk->YdotRHS_slow;
447   Vec             Yslow,Yfast; /* subvectors store the stges of slow components and fast components respectively                           */
448   const PetscInt  s = tab->s;
449   const PetscReal *A = tab->A,*c = tab->c;
450   PetscScalar     *w = rk->work;
451   PetscInt        i,j;
452   PetscReal       next_time_step = ts->time_step,t = ts->ptime,h = ts->time_step;
453   PetscErrorCode  ierr;
454 
455   PetscFunctionBegin;
456   rk->status = TS_STEP_INCOMPLETE;
457   for (i=0; i<s; i++) {
458     ierr = VecGetSubVector(YdotRHS[i],rk->is_slow,&YdotRHS_slow[i]);CHKERRQ(ierr);
459     ierr = VecGetSubVector(YdotRHS[i],rk->is_fast,&YdotRHS_fast[i]);CHKERRQ(ierr);
460   }
461   ierr = VecCopy(ts->vec_sol,rk->X0);CHKERRQ(ierr);
462   /* propogate both slow and fast components using large time steps */
463   for (i=0; i<s; i++) {
464     rk->stage_time = t + h*c[i];
465     ierr = TSPreStage(ts,rk->stage_time);CHKERRQ(ierr);
466     ierr = VecCopy(ts->vec_sol,Y[i]);CHKERRQ(ierr);
467     ierr = VecGetSubVector(Y[i],rk->is_fast,&Yfast);CHKERRQ(ierr);
468     ierr = VecGetSubVector(Y[i],rk->is_slow,&Yslow);CHKERRQ(ierr);
469     for (j=0; j<i; j++) w[j] = h*A[i*s+j];
470     ierr = VecMAXPY(Yfast,i,w,YdotRHS_fast);CHKERRQ(ierr);
471     ierr = VecMAXPY(Yslow,i,w,YdotRHS_slow);CHKERRQ(ierr);
472     ierr = VecRestoreSubVector(Y[i],rk->is_fast,&Yfast);CHKERRQ(ierr);
473     ierr = VecRestoreSubVector(Y[i],rk->is_slow,&Yslow);CHKERRQ(ierr);
474     ierr = TSPostStage(ts,rk->stage_time,i,Y); CHKERRQ(ierr);
475     ierr = TSComputeRHSFunction(rk->subts_slow,t+h*c[i],Y[i],YdotRHS_slow[i]);CHKERRQ(ierr);
476     ierr = TSComputeRHSFunction(rk->subts_fast,t+h*c[i],Y[i],YdotRHS_fast[i]);CHKERRQ(ierr);
477   }
478   rk->slow = PETSC_TRUE;
479   /* update the slow components of the solution using slow time step */
480   ierr = TSEvaluateStep_MRKSPLIT(ts,tab->order,ts->vec_sol,NULL);CHKERRQ(ierr);
481   /* YdotRHS will be used for interpolation during refinement */
482   for (i=0; i<s; i++) {
483     ierr = VecRestoreSubVector(YdotRHS[i],rk->is_slow,&YdotRHS_slow[i]);CHKERRQ(ierr);
484     ierr = VecRestoreSubVector(YdotRHS[i],rk->is_fast,&YdotRHS_fast[i]);CHKERRQ(ierr);
485   }
486 
487   ierr = TSStepRefine_MRKSPLIT(ts);CHKERRQ(ierr);
488 
489   ts->ptime = t + ts->time_step;
490   ts->time_step = next_time_step;
491   rk->status = TS_STEP_COMPLETE;
492   PetscFunctionReturn(0);
493 }
494 
495 /*@C
496   TSRKSetMultirateType - Set the type of RK Multirate scheme
497 
498   Logically collective
499 
500   Input Parameter:
501 +  ts - timestepping context
502 -  mrktype - type of MRK-scheme
503 
504   Options Database:
505 .   -ts_rk_multiarte_type - <none,nonsplit,split>
506 
507   Level: intermediate
508 @*/
509 PetscErrorCode TSRKSetMultirateType(TS ts, TSMRKType mrktype)
510 {
511   TS_RK          *rk = (TS_RK*)ts->data;
512   PetscErrorCode ierr;
513 
514   PetscFunctionBegin;
515   PetscValidHeaderSpecific(ts,TS_CLASSID,1);
516   switch(mrktype){
517     case TSMRKNONE:
518       break;
519     case TSMRKNONSPLIT:
520       ts->ops->step           = TSStep_MRKNONSPLIT;
521       ts->ops->interpolate    = TSInterpolate_MRKNONSPLIT;
522       rk->dtratio             = 2;
523       ierr = TSRKSetType(ts,TSMRKDefault);CHKERRQ(ierr);
524       ierr = PetscObjectComposeFunction((PetscObject)ts,"TSSetUp_MRKNONSPLIT_C",TSSetUp_MRKNONSPLIT);CHKERRQ(ierr);
525       ierr = PetscObjectComposeFunction((PetscObject)ts,"TSReset_MRKNONSPLIT_C",TSReset_MRKNONSPLIT);CHKERRQ(ierr);
526       break;
527     case TSMRKSPLIT:
528       ts->ops->step           = TSStep_MRKSPLIT;
529       ts->ops->evaluatestep   = TSEvaluateStep_MRKSPLIT;
530       ts->ops->interpolate    = TSInterpolate_MRKSPLIT;
531       rk->dtratio             = 2;
532       ierr = TSRKSetType(ts,TSMRKDefault);CHKERRQ(ierr);
533       ierr = PetscObjectComposeFunction((PetscObject)ts,"TSSetUp_MRKSPLIT_C",TSSetUp_MRKSPLIT);CHKERRQ(ierr);
534       ierr = PetscObjectComposeFunction((PetscObject)ts,"TSReset_MRKSPLIT_C",TSReset_MRKSPLIT);CHKERRQ(ierr);
535       break;
536     default :
537       SETERRQ1(PetscObjectComm((PetscObject)ts),PETSC_ERR_ARG_UNKNOWN_TYPE,"Unknown type '%s'",mrktype);
538   }
539   PetscFunctionReturn(0);
540 }
541 
542