xref: /petsc/src/ts/impls/explicit/rk/mrk.c (revision 4b84eec9ecbe260473731d36f19df2cab502356b)
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_MRKSPLIT(TS ts)
21 {
22   TS_RK          *rk = (TS_RK*)ts->data;
23   RKTableau      tab = rk->tableau;
24   DM             dm,subdm,newdm;
25   PetscErrorCode ierr;
26 
27   PetscFunctionBegin;
28   ierr = TSRHSSplitGetIS(ts,"slow",&rk->is_slow);CHKERRQ(ierr);
29   ierr = TSRHSSplitGetIS(ts,"fast",&rk->is_fast);CHKERRQ(ierr);
30   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");
31   ierr = TSRHSSplitGetSubTS(ts,"slow",&rk->subts_slow);CHKERRQ(ierr);
32   ierr = TSRHSSplitGetSubTS(ts,"fast",&rk->subts_fast);CHKERRQ(ierr);
33   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");
34 
35   /* Only copy */
36   ierr = TSGetDM(ts,&dm);CHKERRQ(ierr);
37   ierr = DMClone(dm,&newdm);CHKERRQ(ierr);
38   ierr = TSGetDM(rk->subts_fast,&subdm);CHKERRQ(ierr);
39   ierr = DMCopyDMTS(subdm,newdm);CHKERRQ(ierr);
40   ierr = DMCopyDMSNES(subdm,newdm);CHKERRQ(ierr);
41   ierr = TSSetDM(rk->subts_fast,newdm);CHKERRQ(ierr);
42   ierr = DMDestroy(&newdm);CHKERRQ(ierr);
43   ierr = DMClone(dm,&newdm);CHKERRQ(ierr);
44   ierr = TSGetDM(rk->subts_slow,&subdm);CHKERRQ(ierr);
45   ierr = DMCopyDMTS(subdm,newdm);CHKERRQ(ierr);
46   ierr = DMCopyDMSNES(subdm,newdm);CHKERRQ(ierr);
47   ierr = TSSetDM(rk->subts_slow,newdm);CHKERRQ(ierr);
48   ierr = DMDestroy(&newdm);CHKERRQ(ierr);
49 
50   ierr = PetscMalloc1(tab->s,&rk->YdotRHS_fast);CHKERRQ(ierr);
51   ierr = PetscMalloc1(tab->s,&rk->YdotRHS_slow);CHKERRQ(ierr);
52   ierr = VecDuplicate(ts->vec_sol,&rk->X0);CHKERRQ(ierr);
53   PetscFunctionReturn(0);
54 }
55 
56 static PetscErrorCode TSReset_MRKSPLIT(TS ts)
57 {
58   TS_RK          *rk = (TS_RK*)ts->data;
59   PetscErrorCode ierr;
60 
61   PetscFunctionBegin;
62   ierr = PetscFree(rk->YdotRHS_fast);CHKERRQ(ierr);
63   ierr = PetscFree(rk->YdotRHS_slow);CHKERRQ(ierr);
64   ierr = VecDestroy(&rk->X0);CHKERRQ(ierr);
65   PetscFunctionReturn(0);
66 }
67 
68 static PetscErrorCode TSSetUp_MRKNONSPLIT(TS ts)
69 {
70   TS_RK          *rk = (TS_RK*)ts->data;
71   RKTableau       tab  = rk->tableau;
72   PetscErrorCode ierr;
73 
74   PetscFunctionBegin;
75   ierr = TSRHSSplitGetIS(ts,"slow",&rk->is_slow);CHKERRQ(ierr);
76   ierr = TSRHSSplitGetIS(ts,"fast",&rk->is_fast);CHKERRQ(ierr);
77   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");
78   ierr = VecDuplicate(ts->vec_sol,&rk->X0);CHKERRQ(ierr);
79   ierr = VecDuplicateVecs(ts->vec_sol,tab->s,&rk->YdotRHS_slow);CHKERRQ(ierr);
80   PetscFunctionReturn(0);
81 }
82 
83 static PetscErrorCode TSReset_MRKNONSPLIT(TS ts)
84 {
85   TS_RK          *rk = (TS_RK*)ts->data;
86   RKTableau      tab  = rk->tableau;
87   PetscErrorCode ierr;
88 
89   PetscFunctionBegin;
90   ierr = VecDestroy(&rk->X0);CHKERRQ(ierr);
91   ierr = VecDestroyVecs(tab->s,&rk->YdotRHS_slow);CHKERRQ(ierr);
92   PetscFunctionReturn(0);
93 }
94 
95 static PetscErrorCode TSInterpolate_MRKNONSPLIT(TS ts,PetscReal itime,Vec X)
96 {
97   TS_RK            *rk = (TS_RK*)ts->data;
98   PetscInt         s  = rk->tableau->s,p = rk->tableau->p,i,j;
99   PetscReal        h;
100   PetscReal        tt,t;
101   PetscScalar      *b;
102   const PetscReal  *B = rk->tableau->binterp;
103   PetscErrorCode   ierr;
104 
105   PetscFunctionBegin;
106   if (!B) SETERRQ1(PetscObjectComm((PetscObject)ts),PETSC_ERR_SUP,"TSRK %s does not have an interpolation formula",rk->tableau->name);
107 
108   switch (rk->status) {
109     case TS_STEP_INCOMPLETE:
110     case TS_STEP_PENDING:
111       h = ts->time_step;
112       t = (itime - ts->ptime)/h;
113       break;
114     case TS_STEP_COMPLETE:
115       h = ts->ptime - ts->ptime_prev;
116       t = (itime - ts->ptime)/h + 1; /* In the interval [0,1] */
117       break;
118     default: SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_PLIB,"Invalid TSStepStatus");
119   }
120   ierr = PetscMalloc1(s,&b);CHKERRQ(ierr);
121   for (i=0; i<s; i++) b[i] = 0;
122   for (j=0,tt=t; j<p; j++,tt*=t) {
123     for (i=0; i<s; i++) {
124       b[i]  += h * B[i*p+j] * tt;
125     }
126   }
127   ierr = VecCopy(rk->X0,X);CHKERRQ(ierr);
128   ierr = VecMAXPY(X,s,b,rk->YdotRHS_slow);CHKERRQ(ierr);
129   ierr = PetscFree(b);CHKERRQ(ierr);
130   PetscFunctionReturn(0);
131 }
132 
133 static PetscErrorCode TSInterpolate_MRKSPLIT(TS ts,PetscReal itime,Vec X)
134 {
135   TS_RK            *rk = (TS_RK*)ts->data;
136   PetscInt         s  = rk->tableau->s,p = rk->tableau->p,i,j;
137   Vec              Yslow;    /* vector holds the slow components in Y[0] */
138   PetscReal        h;
139   PetscReal        tt,t;
140   PetscScalar      *b;
141   const PetscReal  *B = rk->tableau->binterp;
142   PetscErrorCode   ierr;
143 
144   PetscFunctionBegin;
145   if (!B) SETERRQ1(PetscObjectComm((PetscObject)ts),PETSC_ERR_SUP,"TSRK %s does not have an interpolation formula",rk->tableau->name);
146 
147   switch (rk->status) {
148     case TS_STEP_INCOMPLETE:
149     case TS_STEP_PENDING:
150       h = ts->time_step;
151       t = (itime - ts->ptime)/h;
152       break;
153     case TS_STEP_COMPLETE:
154       h = ts->ptime - ts->ptime_prev;
155       t = (itime - ts->ptime)/h + 1; /* In the interval [0,1] */
156       break;
157     default: SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_PLIB,"Invalid TSStepStatus");
158   }
159   ierr = PetscMalloc1(s,&b);CHKERRQ(ierr);
160   for (i=0; i<s; i++) b[i] = 0;
161   for (j=0,tt=t; j<p; j++,tt*=t) {
162     for (i=0; i<s; i++) {
163       b[i]  += h * B[i*p+j] * tt;
164     }
165   }
166   ierr = VecGetSubVector(rk->X0,rk->is_slow,&Yslow);CHKERRQ(ierr);
167   ierr = VecCopy(Yslow,X);CHKERRQ(ierr);
168   ierr = VecMAXPY(X,s,b,rk->YdotRHS_slow);CHKERRQ(ierr);
169   ierr = VecRestoreSubVector(rk->X0,rk->is_slow,&Yslow);CHKERRQ(ierr);
170   ierr = PetscFree(b);CHKERRQ(ierr);
171   PetscFunctionReturn(0);
172 }
173 
174 static PetscErrorCode TSStep_MRKNONSPLIT(TS ts)
175 {
176   TS_RK             *rk = (TS_RK*)ts->data;
177   RKTableau         tab  = rk->tableau;
178   Vec               *Y = rk->Y,*YdotRHS = rk->YdotRHS,*YdotRHS_slow = rk->YdotRHS_slow;
179   Vec               stage_slow,sol_slow;   /* vectors store the slow components */
180   Vec               subvec_slow;           /* sub vector to store the slow components */
181   const PetscInt    s = tab->s;
182   const PetscReal   *A = tab->A,*c = tab->c;
183   PetscScalar       *w = rk->work;
184   PetscInt          i,j,k,dtratio = rk->dtratio;
185   PetscReal         next_time_step = ts->time_step,t = ts->ptime,h = ts->time_step;
186   PetscErrorCode    ierr;
187 
188   PetscFunctionBegin;
189   rk->status = TS_STEP_INCOMPLETE;
190   ierr = VecDuplicate(ts->vec_sol,&stage_slow);CHKERRQ(ierr);
191   ierr = VecDuplicate(ts->vec_sol,&sol_slow);CHKERRQ(ierr);
192   ierr = VecCopy(ts->vec_sol,rk->X0);CHKERRQ(ierr);
193   for (i=0; i<s; i++) {
194     rk->stage_time = t + h*c[i];
195     ierr = TSPreStage(ts,rk->stage_time);CHKERRQ(ierr);
196     ierr = VecCopy(ts->vec_sol,Y[i]);CHKERRQ(ierr);
197     for (j=0; j<i; j++) w[j] = h*A[i*s+j];
198     ierr = VecMAXPY(Y[i],i,w,YdotRHS_slow);CHKERRQ(ierr);
199     ierr = TSPostStage(ts,rk->stage_time,i,Y); CHKERRQ(ierr);
200     /* compute the stage RHS */
201     ierr = TSComputeRHSFunction(ts,t+h*c[i],Y[i],YdotRHS_slow[i]);CHKERRQ(ierr);
202   }
203   /* update the slow components in the solution */
204   rk->YdotRHS = YdotRHS_slow;
205   rk->dtratio = 1;
206   ierr = TSEvaluateStep(ts,tab->order,sol_slow,NULL);CHKERRQ(ierr);
207   rk->dtratio = dtratio;
208   rk->YdotRHS = YdotRHS;
209   for (k=0; k<rk->dtratio; k++) {
210     for (i=0; i<s; i++) {
211       rk->stage_time = t + h/rk->dtratio*c[i];
212       ierr = TSPreStage(ts,rk->stage_time);CHKERRQ(ierr);
213       /* 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 */
214       ierr = VecCopy(ts->vec_sol,Y[i]);CHKERRQ(ierr);
215       for (j=0; j<i; j++) w[j] = h/rk->dtratio*A[i*s+j];
216       ierr = VecMAXPY(Y[i],i,w,YdotRHS);CHKERRQ(ierr);
217       ierr = TSInterpolate_MRKNONSPLIT(ts,t+k*h/rk->dtratio+h/rk->dtratio*c[i],stage_slow);CHKERRQ(ierr);
218       /* update the slow components in the stage value */
219       ierr = VecGetSubVector(stage_slow,rk->is_slow,&subvec_slow);CHKERRQ(ierr);
220       ierr = VecISCopy(Y[i],rk->is_slow,SCATTER_FORWARD,subvec_slow);CHKERRQ(ierr);
221       ierr = VecRestoreSubVector(stage_slow,rk->is_slow,&subvec_slow);CHKERRQ(ierr);
222       ierr = TSPostStage(ts,rk->stage_time,i,Y);CHKERRQ(ierr);
223       /* compute the stage RHS */
224       ierr = TSComputeRHSFunction(ts,t+k*h/rk->dtratio+h/rk->dtratio*c[i],Y[i],YdotRHS[i]);CHKERRQ(ierr);
225     }
226     ierr = TSEvaluateStep(ts,tab->order,ts->vec_sol,NULL);CHKERRQ(ierr);
227   }
228   /* update the slow components in the solution */
229   ierr = VecGetSubVector(sol_slow,rk->is_slow,&subvec_slow);CHKERRQ(ierr);
230   ierr = VecISCopy(ts->vec_sol,rk->is_slow,SCATTER_FORWARD,subvec_slow);CHKERRQ(ierr);
231   ierr = VecRestoreSubVector(sol_slow,rk->is_slow,&subvec_slow);CHKERRQ(ierr);
232 
233   ts->ptime += ts->time_step;
234   ts->time_step = next_time_step;
235   rk->status = TS_STEP_COMPLETE;
236   /* free memory of work vectors */
237   ierr = VecDestroy(&stage_slow);CHKERRQ(ierr);
238   ierr = VecDestroy(&sol_slow);CHKERRQ(ierr);
239   PetscFunctionReturn(0);
240 }
241 
242 /*
243  This is for partitioned RHS multirate RK method
244  The step completion formula is
245 
246  x1 = x0 + h b^T YdotRHS
247 
248 */
249 static PetscErrorCode TSEvaluateStep_MRKSPLIT(TS ts,PetscInt order,Vec X,PetscBool *done)
250 {
251   TS_RK           *rk = (TS_RK*)ts->data;
252   RKTableau       tab  = rk->tableau;
253   Vec             Xslow,Xfast;                  /* subvectors of X which store slow components and fast components respectively */
254   PetscScalar     *w = rk->work;
255   PetscReal       h = ts->time_step;
256   PetscInt        s = tab->s,j;
257   PetscErrorCode  ierr;
258 
259   PetscFunctionBegin;
260   ierr = VecCopy(ts->vec_sol,X);CHKERRQ(ierr);
261   if (rk->slow) {
262     for (j=0; j<s; j++) w[j] = h*tab->b[j];
263     ierr = VecGetSubVector(ts->vec_sol,rk->is_slow,&Xslow);CHKERRQ(ierr);
264     ierr = VecMAXPY(Xslow,s,w,rk->YdotRHS_slow);CHKERRQ(ierr);
265     ierr = VecRestoreSubVector(ts->vec_sol,rk->is_slow,&Xslow);CHKERRQ(ierr);;
266   } else {
267     for (j=0; j<s; j++) w[j] = h/rk->dtratio*tab->b[j];
268     ierr = VecGetSubVector(X,rk->is_fast,&Xfast);CHKERRQ(ierr);
269     ierr = VecMAXPY(Xfast,s,w,rk->YdotRHS_fast);CHKERRQ(ierr);
270     ierr = VecRestoreSubVector(X,rk->is_fast,&Xfast);CHKERRQ(ierr);
271   }
272   PetscFunctionReturn(0);
273 }
274 
275 static PetscErrorCode TSStep_MRKSPLIT(TS ts)
276 {
277   TS_RK             *rk = (TS_RK*)ts->data;
278   RKTableau         tab = rk->tableau;
279   Vec               *Y = rk->Y,*YdotRHS = rk->YdotRHS;
280   Vec               *YdotRHS_fast = rk->YdotRHS_fast,*YdotRHS_slow = rk->YdotRHS_slow;
281   Vec               Yslow,Yfast;                         /* subvectors store the stges of slow components and fast components respectively                           */
282   const PetscInt    s = tab->s;
283   const PetscReal   *A = tab->A,*c = tab->c;
284   PetscScalar       *w = rk->work;
285   PetscInt          i,j,k;
286   PetscReal         next_time_step = ts->time_step,t = ts->ptime,h = ts->time_step;
287   PetscErrorCode    ierr;
288 
289   PetscFunctionBegin;
290   rk->status = TS_STEP_INCOMPLETE;
291   for (i=0; i<s; i++) {
292     ierr = VecGetSubVector(YdotRHS[i],rk->is_slow,&YdotRHS_slow[i]);CHKERRQ(ierr);
293     ierr = VecGetSubVector(YdotRHS[i],rk->is_fast,&YdotRHS_fast[i]);CHKERRQ(ierr);
294   }
295   ierr = VecCopy(ts->vec_sol,rk->X0);CHKERRQ(ierr);
296   /* propogate both slow and fast components using large time steps */
297   for (i=0; i<s; i++) {
298     rk->stage_time = t + h*c[i];
299     ierr = TSPreStage(ts,rk->stage_time);CHKERRQ(ierr);
300     ierr = VecCopy(ts->vec_sol,Y[i]);CHKERRQ(ierr);
301     /*ierr = VecCopy(ts->vec_sol,Yc[i]);CHKERRQ(ierr);*/
302     ierr = VecGetSubVector(Y[i],rk->is_slow,&Yslow);CHKERRQ(ierr);
303     ierr = VecGetSubVector(Y[i],rk->is_fast,&Yfast);CHKERRQ(ierr);
304     for (j=0; j<i; j++) w[j] = h*A[i*s+j];
305     ierr = VecMAXPY(Yslow,i,w,YdotRHS_slow);CHKERRQ(ierr);
306     ierr = VecMAXPY(Yfast,i,w,YdotRHS_fast);CHKERRQ(ierr);
307     ierr = VecRestoreSubVector(Y[i],rk->is_slow,&Yslow);CHKERRQ(ierr);
308     ierr = VecRestoreSubVector(Y[i],rk->is_fast,&Yfast);CHKERRQ(ierr);
309     ierr = TSPostStage(ts,rk->stage_time,i,Y); CHKERRQ(ierr);
310     ierr = TSComputeRHSFunction(rk->subts_slow,t+h*c[i],Y[i],YdotRHS_slow[i]);CHKERRQ(ierr);
311     ierr = TSComputeRHSFunction(rk->subts_fast,t+h*c[i],Y[i],YdotRHS_fast[i]);CHKERRQ(ierr);
312   }
313   rk->slow = PETSC_TRUE;
314   ierr = TSEvaluateStep_MRKSPLIT(ts,tab->order,ts->vec_sol,NULL);CHKERRQ(ierr);
315   for (k=0; k<rk->dtratio; k++) {
316     /* propogate fast component using small time steps */
317     for (i=0; i<s; i++) {
318       rk->stage_time = t + h/rk->dtratio*c[i];
319       ierr = TSPreStage(ts,rk->stage_time);CHKERRQ(ierr);
320       ierr = VecCopy(ts->vec_sol,Y[i]);CHKERRQ(ierr);
321       /* stage value for fast components */
322       for (j=0; j<i; j++) w[j] = h/rk->dtratio*A[i*s+j];
323       ierr = VecGetSubVector(Y[i],rk->is_fast,&Yfast);CHKERRQ(ierr);
324       ierr = VecMAXPY(Yfast,i,w,YdotRHS_fast);CHKERRQ(ierr);
325       ierr = VecRestoreSubVector(Y[i],rk->is_fast,&Yfast);CHKERRQ(ierr);
326       /* stage value for slow components */
327       ierr = VecGetSubVector(Y[i],rk->is_slow,&Yslow);CHKERRQ(ierr);
328       ierr = TSInterpolate_MRKSPLIT(ts,t+k*h/rk->dtratio+h/rk->dtratio*c[i],Yslow);CHKERRQ(ierr);
329       ierr = VecRestoreSubVector(Y[i],rk->is_slow,&Yslow);CHKERRQ(ierr);
330       ierr = TSPostStage(ts,rk->stage_time,i,Y); CHKERRQ(ierr);
331       /* compute the stage RHS for fast components */
332       ierr = TSComputeRHSFunction(rk->subts_fast,t+k*h/rk->dtratio+h/rk->dtratio*c[i],Y[i],YdotRHS_fast[i]);CHKERRQ(ierr);
333     }
334     /* update the value of fast components when using fast time step */
335     rk->slow = PETSC_FALSE;
336     ierr = TSEvaluateStep_MRKSPLIT(ts,tab->order,ts->vec_sol,NULL);CHKERRQ(ierr);
337   }
338   for (i=0; i<s; i++) {
339     ierr = VecRestoreSubVector(YdotRHS[i],rk->is_slow,&YdotRHS_slow[i]);CHKERRQ(ierr);
340     ierr = VecRestoreSubVector(YdotRHS[i],rk->is_fast,&YdotRHS_fast[i]);CHKERRQ(ierr);
341   }
342   ts->ptime += ts->time_step;
343   ts->time_step = next_time_step;
344   rk->status = TS_STEP_COMPLETE;
345   PetscFunctionReturn(0);
346 }
347 
348 /*@C
349   TSRKSetMultirateType - Set the type of RK Multirate scheme
350 
351   Logically collective
352 
353   Input Parameter:
354 +  ts - timestepping context
355 -  mrktype - type of MRK-scheme
356 
357   Options Database:
358 .   -ts_rk_multiarte_type - <none,nonsplit,split>
359 
360   Level: intermediate
361 @*/
362 PetscErrorCode TSRKSetMultirateType(TS ts, TSMRKType mrktype)
363 {
364   TS_RK          *rk = (TS_RK*)ts->data;
365   PetscErrorCode ierr;
366 
367   PetscFunctionBegin;
368   PetscValidHeaderSpecific(ts,TS_CLASSID,1);
369   switch(mrktype){
370     case TSMRKNONE:
371       break;
372     case TSMRKNONSPLIT:
373       ts->ops->step           = TSStep_MRKNONSPLIT;
374       ts->ops->interpolate    = TSInterpolate_MRKNONSPLIT;
375       rk->dtratio             = 2;
376       ierr = TSRKSetType(ts,TSMRKDefault);CHKERRQ(ierr);
377       ierr = PetscObjectComposeFunction((PetscObject)ts,"TSSetUp_MRKNONSPLIT_C",TSSetUp_MRKNONSPLIT);CHKERRQ(ierr);
378       ierr = PetscObjectComposeFunction((PetscObject)ts,"TSReset_MRKNONSPLIT_C",TSReset_MRKNONSPLIT);CHKERRQ(ierr);
379       break;
380     case TSMRKSPLIT:
381       ts->ops->step           = TSStep_MRKSPLIT;
382       ts->ops->evaluatestep   = TSEvaluateStep_MRKSPLIT;
383       ts->ops->interpolate    = TSInterpolate_MRKSPLIT;
384       rk->dtratio             = 2;
385       ierr = TSRKSetType(ts,TSMRKDefault);CHKERRQ(ierr);
386       ierr = PetscObjectComposeFunction((PetscObject)ts,"TSSetUp_MRKSPLIT_C",TSSetUp_MRKSPLIT);CHKERRQ(ierr);
387       ierr = PetscObjectComposeFunction((PetscObject)ts,"TSReset_MRKSPLIT_C",TSReset_MRKSPLIT);CHKERRQ(ierr);
388       break;
389     default :
390       SETERRQ1(PetscObjectComm((PetscObject)ts),PETSC_ERR_ARG_UNKNOWN_TYPE,"Unknown type '%s'",mrktype);
391   }
392   PetscFunctionReturn(0);
393 }
394 
395