xref: /petsc/src/ts/impls/explicit/rk/mrk.c (revision 9566063d113dddea24716c546802770db7481bc0)
1474dd773SHong Zhang /*
2474dd773SHong Zhang   Code for time stepping with the multi-rate Runge-Kutta method
3474dd773SHong Zhang 
4474dd773SHong Zhang   Notes:
5474dd773SHong Zhang   1) The general system is written as
6474dd773SHong Zhang      Udot = F(t,U) for the nonsplit version of multi-rate RK,
7474dd773SHong Zhang      user should give the indexes for both slow and fast components;
8474dd773SHong Zhang   2) The general system is written as
9474dd773SHong Zhang      Usdot = Fs(t,Us,Uf)
10474dd773SHong Zhang      Ufdot = Ff(t,Us,Uf) for multi-rate RK with RHS splits,
11474dd773SHong Zhang      user should partioned RHS by themselves and also provide the indexes for both slow and fast components.
12474dd773SHong Zhang */
13474dd773SHong Zhang 
14474dd773SHong Zhang #include <petsc/private/tsimpl.h>
15474dd773SHong Zhang #include <petscdm.h>
16474dd773SHong Zhang #include <../src/ts/impls/explicit/rk/rk.h>
1721052c0cSHong Zhang #include <../src/ts/impls/explicit/rk/mrk.h>
18474dd773SHong Zhang 
1902555c94SHong Zhang static PetscErrorCode TSReset_RK_MultirateNonsplit(TS ts)
20e5bffa22SHong Zhang {
21e5bffa22SHong Zhang   TS_RK          *rk = (TS_RK*)ts->data;
22e5bffa22SHong Zhang   RKTableau      tab = rk->tableau;
23e5bffa22SHong Zhang 
24e5bffa22SHong Zhang   PetscFunctionBegin;
25*9566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&rk->X0));
26*9566063dSJacob Faibussowitsch   PetscCall(VecDestroyVecs(tab->s,&rk->YdotRHS_slow));
27e5bffa22SHong Zhang   PetscFunctionReturn(0);
28e5bffa22SHong Zhang }
29e5bffa22SHong Zhang 
3002555c94SHong Zhang static PetscErrorCode TSInterpolate_RK_MultirateNonsplit(TS ts,PetscReal itime,Vec X)
31e5bffa22SHong Zhang {
32e5bffa22SHong Zhang   TS_RK            *rk = (TS_RK*)ts->data;
33e5bffa22SHong Zhang   PetscInt         s = rk->tableau->s,p = rk->tableau->p,i,j;
3463a6f1b4SHong Zhang   PetscReal        h = ts->time_step;
35e5bffa22SHong Zhang   PetscReal        tt,t;
36e5bffa22SHong Zhang   PetscScalar      *b;
37e5bffa22SHong Zhang   const PetscReal  *B = rk->tableau->binterp;
38e5bffa22SHong Zhang 
39e5bffa22SHong Zhang   PetscFunctionBegin;
403c633725SBarry Smith   PetscCheck(B,PetscObjectComm((PetscObject)ts),PETSC_ERR_SUP,"TSRK %s does not have an interpolation formula",rk->tableau->name);
4163a6f1b4SHong Zhang   t = (itime - rk->ptime)/h;
42*9566063dSJacob Faibussowitsch   PetscCall(PetscMalloc1(s,&b));
43e5bffa22SHong Zhang   for (i=0; i<s; i++) b[i] = 0;
44e5bffa22SHong Zhang   for (j=0,tt=t; j<p; j++,tt*=t) {
45e5bffa22SHong Zhang     for (i=0; i<s; i++) {
46e5bffa22SHong Zhang       b[i]  += h * B[i*p+j] * tt;
47e5bffa22SHong Zhang     }
48e5bffa22SHong Zhang   }
49*9566063dSJacob Faibussowitsch   PetscCall(VecCopy(rk->X0,X));
50*9566063dSJacob Faibussowitsch   PetscCall(VecMAXPY(X,s,b,rk->YdotRHS_slow));
51*9566063dSJacob Faibussowitsch   PetscCall(PetscFree(b));
52e5bffa22SHong Zhang   PetscFunctionReturn(0);
53e5bffa22SHong Zhang }
54e5bffa22SHong Zhang 
5502555c94SHong Zhang static PetscErrorCode TSStepRefine_RK_MultirateNonsplit(TS ts)
5663a6f1b4SHong Zhang {
5763a6f1b4SHong Zhang   TS              previousts,subts;
5863a6f1b4SHong Zhang   TS_RK           *rk = (TS_RK*)ts->data;
5963a6f1b4SHong Zhang   RKTableau       tab = rk->tableau;
6063a6f1b4SHong Zhang   Vec             *Y = rk->Y,*YdotRHS = rk->YdotRHS;
6163a6f1b4SHong Zhang   Vec             vec_fast,subvec_fast;
6263a6f1b4SHong Zhang   const PetscInt  s = tab->s;
6363a6f1b4SHong Zhang   const PetscReal *A = tab->A,*c = tab->c;
6463a6f1b4SHong Zhang   PetscScalar     *w = rk->work;
6563a6f1b4SHong Zhang   PetscInt        i,j,k;
6663a6f1b4SHong Zhang   PetscReal       t = ts->ptime,h = ts->time_step;
6763a6f1b4SHong Zhang 
68362febeeSStefano Zampini   PetscFunctionBegin;
69*9566063dSJacob Faibussowitsch   PetscCall(VecDuplicate(ts->vec_sol,&vec_fast));
7063a6f1b4SHong Zhang   previousts = rk->subts_current;
71*9566063dSJacob Faibussowitsch   PetscCall(TSRHSSplitGetSubTS(rk->subts_current,"fast",&subts));
72*9566063dSJacob Faibussowitsch   PetscCall(TSRHSSplitGetSubTS(subts,"fast",&subts));
7363a6f1b4SHong Zhang   for (k=0; k<rk->dtratio; k++) {
7463a6f1b4SHong Zhang     for (i=0; i<s; i++) {
75*9566063dSJacob Faibussowitsch       PetscCall(TSInterpolate_RK_MultirateNonsplit(ts,t+k*h/rk->dtratio+h/rk->dtratio*c[i],Y[i]));
7663a6f1b4SHong Zhang       for (j=0; j<i; j++) w[j] = h/rk->dtratio*A[i*s+j];
7763a6f1b4SHong Zhang       /* 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 */
78*9566063dSJacob Faibussowitsch       PetscCall(VecCopy(ts->vec_sol,vec_fast));
79*9566063dSJacob Faibussowitsch       PetscCall(VecMAXPY(vec_fast,i,w,YdotRHS));
8063a6f1b4SHong Zhang       /* update the fast components in the stage value */
81*9566063dSJacob Faibussowitsch       PetscCall(VecGetSubVector(vec_fast,rk->is_fast,&subvec_fast));
82*9566063dSJacob Faibussowitsch       PetscCall(VecISCopy(Y[i],rk->is_fast,SCATTER_FORWARD,subvec_fast));
83*9566063dSJacob Faibussowitsch       PetscCall(VecRestoreSubVector(vec_fast,rk->is_fast,&subvec_fast));
8463a6f1b4SHong Zhang       /* compute the stage RHS */
85*9566063dSJacob Faibussowitsch       PetscCall(TSComputeRHSFunction(ts,t+k*h/rk->dtratio+h/rk->dtratio*c[i],Y[i],YdotRHS[i]));
8663a6f1b4SHong Zhang     }
87*9566063dSJacob Faibussowitsch     PetscCall(VecCopy(ts->vec_sol,vec_fast));
88*9566063dSJacob Faibussowitsch     PetscCall(TSEvaluateStep(ts,tab->order,vec_fast,NULL));
8963a6f1b4SHong Zhang     /* update the fast components in the solution */
90*9566063dSJacob Faibussowitsch     PetscCall(VecGetSubVector(vec_fast,rk->is_fast,&subvec_fast));
91*9566063dSJacob Faibussowitsch     PetscCall(VecISCopy(ts->vec_sol,rk->is_fast,SCATTER_FORWARD,subvec_fast));
92*9566063dSJacob Faibussowitsch     PetscCall(VecRestoreSubVector(vec_fast,rk->is_fast,&subvec_fast));
9363a6f1b4SHong Zhang 
9463a6f1b4SHong Zhang     if (subts) {
9563a6f1b4SHong Zhang       Vec *YdotRHS_copy;
96*9566063dSJacob Faibussowitsch       PetscCall(VecDuplicateVecs(ts->vec_sol,s,&YdotRHS_copy));
9763a6f1b4SHong Zhang       rk->subts_current = rk->subts_fast;
9863a6f1b4SHong Zhang       ts->ptime = t+k*h/rk->dtratio;
9963a6f1b4SHong Zhang       ts->time_step = h/rk->dtratio;
100*9566063dSJacob Faibussowitsch       PetscCall(TSRHSSplitGetIS(rk->subts_current,"fast",&rk->is_fast));
10163a6f1b4SHong Zhang       for (i=0; i<s; i++) {
102*9566063dSJacob Faibussowitsch         PetscCall(VecCopy(rk->YdotRHS_slow[i],YdotRHS_copy[i]));
103*9566063dSJacob Faibussowitsch         PetscCall(VecCopy(YdotRHS[i],rk->YdotRHS_slow[i]));
10463a6f1b4SHong Zhang       }
10563a6f1b4SHong Zhang 
106*9566063dSJacob Faibussowitsch       PetscCall(TSStepRefine_RK_MultirateNonsplit(ts));
10763a6f1b4SHong Zhang 
10863a6f1b4SHong Zhang       rk->subts_current = previousts;
10963a6f1b4SHong Zhang       ts->ptime = t;
11063a6f1b4SHong Zhang       ts->time_step = h;
111*9566063dSJacob Faibussowitsch       PetscCall(TSRHSSplitGetIS(previousts,"fast",&rk->is_fast));
11263a6f1b4SHong Zhang       for (i=0; i<s; i++) {
113*9566063dSJacob Faibussowitsch         PetscCall(VecCopy(YdotRHS_copy[i],rk->YdotRHS_slow[i]));
11463a6f1b4SHong Zhang       }
115*9566063dSJacob Faibussowitsch       PetscCall(VecDestroyVecs(s,&YdotRHS_copy));
11663a6f1b4SHong Zhang     }
11763a6f1b4SHong Zhang   }
118*9566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&vec_fast));
11963a6f1b4SHong Zhang   PetscFunctionReturn(0);
12063a6f1b4SHong Zhang }
12163a6f1b4SHong Zhang 
12202555c94SHong Zhang static PetscErrorCode TSStep_RK_MultirateNonsplit(TS ts)
123e5bffa22SHong Zhang {
124e5bffa22SHong Zhang   TS_RK           *rk = (TS_RK*)ts->data;
125e5bffa22SHong Zhang   RKTableau       tab = rk->tableau;
126e5bffa22SHong Zhang   Vec             *Y = rk->Y,*YdotRHS = rk->YdotRHS,*YdotRHS_slow = rk->YdotRHS_slow;
127e5bffa22SHong Zhang   Vec             stage_slow,sol_slow; /* vectors store the slow components */
128e5bffa22SHong Zhang   Vec             subvec_slow; /* sub vector to store the slow components */
129e5bffa22SHong Zhang   IS              is_slow = rk->is_slow;
130e5bffa22SHong Zhang   const PetscInt  s = tab->s;
131e5bffa22SHong Zhang   const PetscReal *A = tab->A,*c = tab->c;
132e5bffa22SHong Zhang   PetscScalar     *w = rk->work;
13363a6f1b4SHong Zhang   PetscInt        i,j,dtratio = rk->dtratio;
134e5bffa22SHong Zhang   PetscReal       next_time_step = ts->time_step,t = ts->ptime,h = ts->time_step;
135e5bffa22SHong Zhang 
136e5bffa22SHong Zhang   PetscFunctionBegin;
137e5bffa22SHong Zhang   rk->status = TS_STEP_INCOMPLETE;
138*9566063dSJacob Faibussowitsch   PetscCall(VecDuplicate(ts->vec_sol,&stage_slow));
139*9566063dSJacob Faibussowitsch   PetscCall(VecDuplicate(ts->vec_sol,&sol_slow));
140*9566063dSJacob Faibussowitsch   PetscCall(VecCopy(ts->vec_sol,rk->X0));
141e5bffa22SHong Zhang   for (i=0; i<s; i++) {
142e5bffa22SHong Zhang     rk->stage_time = t + h*c[i];
143*9566063dSJacob Faibussowitsch     PetscCall(TSPreStage(ts,rk->stage_time));
144*9566063dSJacob Faibussowitsch     PetscCall(VecCopy(ts->vec_sol,Y[i]));
145e5bffa22SHong Zhang     for (j=0; j<i; j++) w[j] = h*A[i*s+j];
146*9566063dSJacob Faibussowitsch     PetscCall(VecMAXPY(Y[i],i,w,YdotRHS_slow));
147*9566063dSJacob Faibussowitsch     PetscCall(TSPostStage(ts,rk->stage_time,i,Y));
148e5bffa22SHong Zhang     /* compute the stage RHS */
149*9566063dSJacob Faibussowitsch     PetscCall(TSComputeRHSFunction(ts,t+h*c[i],Y[i],YdotRHS_slow[i]));
150e5bffa22SHong Zhang   }
151e5bffa22SHong Zhang   /* update the slow components in the solution */
152e5bffa22SHong Zhang   rk->YdotRHS = YdotRHS_slow;
153e5bffa22SHong Zhang   rk->dtratio = 1;
154*9566063dSJacob Faibussowitsch   PetscCall(TSEvaluateStep(ts,tab->order,sol_slow,NULL));
155e5bffa22SHong Zhang   rk->dtratio = dtratio;
156e5bffa22SHong Zhang   rk->YdotRHS = YdotRHS;
157e5bffa22SHong Zhang   /* update the slow components in the solution */
158*9566063dSJacob Faibussowitsch   PetscCall(VecGetSubVector(sol_slow,is_slow,&subvec_slow));
159*9566063dSJacob Faibussowitsch   PetscCall(VecISCopy(ts->vec_sol,is_slow,SCATTER_FORWARD,subvec_slow));
160*9566063dSJacob Faibussowitsch   PetscCall(VecRestoreSubVector(sol_slow,is_slow,&subvec_slow));
161e5bffa22SHong Zhang 
16263a6f1b4SHong Zhang   rk->subts_current = ts;
16363a6f1b4SHong Zhang   rk->ptime = t;
16463a6f1b4SHong Zhang   rk->time_step = h;
165*9566063dSJacob Faibussowitsch   PetscCall(TSStepRefine_RK_MultirateNonsplit(ts));
16663a6f1b4SHong Zhang 
167e5bffa22SHong Zhang   ts->ptime = t + ts->time_step;
168e5bffa22SHong Zhang   ts->time_step = next_time_step;
169e5bffa22SHong Zhang   rk->status = TS_STEP_COMPLETE;
17063a6f1b4SHong Zhang 
171e5bffa22SHong Zhang   /* free memory of work vectors */
172*9566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&stage_slow));
173*9566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&sol_slow));
174e5bffa22SHong Zhang   PetscFunctionReturn(0);
175e5bffa22SHong Zhang }
176e5bffa22SHong Zhang 
17702555c94SHong Zhang static PetscErrorCode TSSetUp_RK_MultirateNonsplit(TS ts)
1780fe4d17eSHong Zhang {
1790fe4d17eSHong Zhang   TS_RK          *rk = (TS_RK*)ts->data;
1800fe4d17eSHong Zhang   RKTableau      tab = rk->tableau;
1810fe4d17eSHong Zhang 
1820fe4d17eSHong Zhang   PetscFunctionBegin;
183*9566063dSJacob Faibussowitsch   PetscCall(TSRHSSplitGetIS(ts,"slow",&rk->is_slow));
184*9566063dSJacob Faibussowitsch   PetscCall(TSRHSSplitGetIS(ts,"fast",&rk->is_fast));
1853c633725SBarry Smith   PetscCheck(rk->is_slow && rk->is_fast,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");
186*9566063dSJacob Faibussowitsch   PetscCall(TSRHSSplitGetSubTS(ts,"slow",&rk->subts_slow));
187*9566063dSJacob Faibussowitsch   PetscCall(TSRHSSplitGetSubTS(ts,"fast",&rk->subts_fast));
1883c633725SBarry Smith   PetscCheck(rk->subts_slow && rk->subts_fast,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");
189*9566063dSJacob Faibussowitsch   PetscCall(VecDuplicate(ts->vec_sol,&rk->X0));
190*9566063dSJacob Faibussowitsch   PetscCall(VecDuplicateVecs(ts->vec_sol,tab->s,&rk->YdotRHS_slow));
1910fe4d17eSHong Zhang   rk->subts_current = rk->subts_fast;
1920fe4d17eSHong Zhang 
19302555c94SHong Zhang   ts->ops->step        = TSStep_RK_MultirateNonsplit;
19402555c94SHong Zhang   ts->ops->interpolate = TSInterpolate_RK_MultirateNonsplit;
1950fe4d17eSHong Zhang   PetscFunctionReturn(0);
1960fe4d17eSHong Zhang }
1970fe4d17eSHong Zhang 
19863a6f1b4SHong Zhang /*
19963a6f1b4SHong Zhang   Copy DM from tssrc to tsdest, while keeping the original DMTS and DMSNES in tsdest.
20063a6f1b4SHong Zhang */
20163a6f1b4SHong Zhang static PetscErrorCode TSCopyDM(TS tssrc,TS tsdest)
20263a6f1b4SHong Zhang {
20363a6f1b4SHong Zhang   DM             newdm,dmsrc,dmdest;
20463a6f1b4SHong Zhang 
20563a6f1b4SHong Zhang   PetscFunctionBegin;
206*9566063dSJacob Faibussowitsch   PetscCall(TSGetDM(tssrc,&dmsrc));
207*9566063dSJacob Faibussowitsch   PetscCall(DMClone(dmsrc,&newdm));
208*9566063dSJacob Faibussowitsch   PetscCall(TSGetDM(tsdest,&dmdest));
209*9566063dSJacob Faibussowitsch   PetscCall(DMCopyDMTS(dmdest,newdm));
210*9566063dSJacob Faibussowitsch   PetscCall(DMCopyDMSNES(dmdest,newdm));
211*9566063dSJacob Faibussowitsch   PetscCall(TSSetDM(tsdest,newdm));
212*9566063dSJacob Faibussowitsch   PetscCall(DMDestroy(&newdm));
21363a6f1b4SHong Zhang   PetscFunctionReturn(0);
21463a6f1b4SHong Zhang }
21563a6f1b4SHong Zhang 
21602555c94SHong Zhang static PetscErrorCode TSReset_RK_MultirateSplit(TS ts)
217474dd773SHong Zhang {
218474dd773SHong Zhang   TS_RK          *rk = (TS_RK*)ts->data;
219474dd773SHong Zhang 
220474dd773SHong Zhang   PetscFunctionBegin;
221bb42530cSHong Zhang   if (rk->subts_slow) {
222*9566063dSJacob Faibussowitsch     PetscCall(PetscFree(rk->subts_slow->data));
223bb42530cSHong Zhang     rk->subts_slow = NULL;
224bb42530cSHong Zhang   }
225bb42530cSHong Zhang   if (rk->subts_fast) {
226*9566063dSJacob Faibussowitsch     PetscCall(PetscFree(rk->YdotRHS_fast));
227*9566063dSJacob Faibussowitsch     PetscCall(PetscFree(rk->YdotRHS_slow));
228*9566063dSJacob Faibussowitsch     PetscCall(VecDestroy(&rk->X0));
229*9566063dSJacob Faibussowitsch     PetscCall(TSReset_RK_MultirateSplit(rk->subts_fast));
230*9566063dSJacob Faibussowitsch     PetscCall(PetscFree(rk->subts_fast->data));
231bb42530cSHong Zhang     rk->subts_fast = NULL;
232bb42530cSHong Zhang   }
233474dd773SHong Zhang   PetscFunctionReturn(0);
234474dd773SHong Zhang }
235474dd773SHong Zhang 
23602555c94SHong Zhang static PetscErrorCode TSInterpolate_RK_MultirateSplit(TS ts,PetscReal itime,Vec X)
237474dd773SHong Zhang {
238474dd773SHong Zhang   TS_RK           *rk = (TS_RK*)ts->data;
23963a6f1b4SHong Zhang   Vec             Xslow;
240474dd773SHong Zhang   PetscInt        s = rk->tableau->s,p = rk->tableau->p,i,j;
241474dd773SHong Zhang   PetscReal       h;
242474dd773SHong Zhang   PetscReal       tt,t;
243474dd773SHong Zhang   PetscScalar     *b;
244474dd773SHong Zhang   const PetscReal *B = rk->tableau->binterp;
245474dd773SHong Zhang 
246474dd773SHong Zhang   PetscFunctionBegin;
2473c633725SBarry Smith   PetscCheck(B,PetscObjectComm((PetscObject)ts),PETSC_ERR_SUP,"TSRK %s does not have an interpolation formula",rk->tableau->name);
248474dd773SHong Zhang 
249474dd773SHong Zhang   switch (rk->status) {
250474dd773SHong Zhang     case TS_STEP_INCOMPLETE:
251474dd773SHong Zhang     case TS_STEP_PENDING:
252474dd773SHong Zhang       h = ts->time_step;
253474dd773SHong Zhang       t = (itime - ts->ptime)/h;
254474dd773SHong Zhang       break;
255474dd773SHong Zhang     case TS_STEP_COMPLETE:
256474dd773SHong Zhang       h = ts->ptime - ts->ptime_prev;
257474dd773SHong Zhang       t = (itime - ts->ptime)/h + 1; /* In the interval [0,1] */
258474dd773SHong Zhang       break;
259474dd773SHong Zhang     default: SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_PLIB,"Invalid TSStepStatus");
260474dd773SHong Zhang   }
261*9566063dSJacob Faibussowitsch   PetscCall(PetscMalloc1(s,&b));
262474dd773SHong Zhang   for (i=0; i<s; i++) b[i] = 0;
263474dd773SHong Zhang   for (j=0,tt=t; j<p; j++,tt*=t) {
264474dd773SHong Zhang     for (i=0; i<s; i++) {
265474dd773SHong Zhang       b[i]  += h * B[i*p+j] * tt;
266474dd773SHong Zhang     }
267474dd773SHong Zhang   }
26863a6f1b4SHong Zhang   for (i=0; i<s; i++) {
269*9566063dSJacob Faibussowitsch     PetscCall(VecGetSubVector(rk->YdotRHS[i],rk->is_slow,&rk->YdotRHS_slow[i]));
27063a6f1b4SHong Zhang   }
271*9566063dSJacob Faibussowitsch   PetscCall(VecGetSubVector(X,rk->is_slow,&Xslow));
272*9566063dSJacob Faibussowitsch   PetscCall(VecISCopy(rk->X0,rk->is_slow,SCATTER_REVERSE,Xslow));
273*9566063dSJacob Faibussowitsch   PetscCall(VecMAXPY(Xslow,s,b,rk->YdotRHS_slow));
274*9566063dSJacob Faibussowitsch   PetscCall(VecRestoreSubVector(X,rk->is_slow,&Xslow));
27563a6f1b4SHong Zhang   for (i=0; i<s; i++) {
276*9566063dSJacob Faibussowitsch     PetscCall(VecRestoreSubVector(rk->YdotRHS[i],rk->is_slow,&rk->YdotRHS_slow[i]));
27763a6f1b4SHong Zhang   }
278*9566063dSJacob Faibussowitsch   PetscCall(PetscFree(b));
279474dd773SHong Zhang   PetscFunctionReturn(0);
280474dd773SHong Zhang }
281474dd773SHong Zhang 
282474dd773SHong Zhang /*
283474dd773SHong Zhang  This is for partitioned RHS multirate RK method
284474dd773SHong Zhang  The step completion formula is
285474dd773SHong Zhang 
286474dd773SHong Zhang  x1 = x0 + h b^T YdotRHS
287474dd773SHong Zhang 
288474dd773SHong Zhang */
28902555c94SHong Zhang static PetscErrorCode TSEvaluateStep_RK_MultirateSplit(TS ts,PetscInt order,Vec X,PetscBool *done)
290474dd773SHong Zhang {
291474dd773SHong Zhang   TS_RK          *rk = (TS_RK*)ts->data;
292474dd773SHong Zhang   RKTableau      tab = rk->tableau;
293474dd773SHong Zhang   Vec            Xslow,Xfast;                  /* subvectors of X which store slow components and fast components respectively */
294474dd773SHong Zhang   PetscScalar    *w = rk->work;
295474dd773SHong Zhang   PetscReal      h = ts->time_step;
296474dd773SHong Zhang   PetscInt       s = tab->s,j;
297474dd773SHong Zhang 
298474dd773SHong Zhang   PetscFunctionBegin;
299*9566063dSJacob Faibussowitsch   PetscCall(VecCopy(ts->vec_sol,X));
300474dd773SHong Zhang   if (rk->slow) {
301474dd773SHong Zhang     for (j=0; j<s; j++) w[j] = h*tab->b[j];
302*9566063dSJacob Faibussowitsch     PetscCall(VecGetSubVector(ts->vec_sol,rk->is_slow,&Xslow));
303*9566063dSJacob Faibussowitsch     PetscCall(VecMAXPY(Xslow,s,w,rk->YdotRHS_slow));
304*9566063dSJacob Faibussowitsch     PetscCall(VecRestoreSubVector(ts->vec_sol,rk->is_slow,&Xslow));
305474dd773SHong Zhang   } else {
306474dd773SHong Zhang     for (j=0; j<s; j++) w[j] = h/rk->dtratio*tab->b[j];
307*9566063dSJacob Faibussowitsch     PetscCall(VecGetSubVector(X,rk->is_fast,&Xfast));
308*9566063dSJacob Faibussowitsch     PetscCall(VecMAXPY(Xfast,s,w,rk->YdotRHS_fast));
309*9566063dSJacob Faibussowitsch     PetscCall(VecRestoreSubVector(X,rk->is_fast,&Xfast));
310474dd773SHong Zhang   }
311474dd773SHong Zhang   PetscFunctionReturn(0);
312474dd773SHong Zhang }
313474dd773SHong Zhang 
31402555c94SHong Zhang static PetscErrorCode TSStepRefine_RK_MultirateSplit(TS ts)
31563a6f1b4SHong Zhang {
31663a6f1b4SHong Zhang   TS_RK           *rk = (TS_RK*)ts->data;
31763a6f1b4SHong Zhang   TS              subts_fast = rk->subts_fast,currentlevelts;
31863a6f1b4SHong Zhang   TS_RK           *subrk_fast = (TS_RK*)subts_fast->data;
31963a6f1b4SHong Zhang   RKTableau       tab = rk->tableau;
32063a6f1b4SHong Zhang   Vec             *Y = rk->Y;
32163a6f1b4SHong Zhang   Vec             *YdotRHS = rk->YdotRHS,*YdotRHS_fast = rk->YdotRHS_fast;
32263a6f1b4SHong Zhang   Vec             Yfast,Xfast;
32363a6f1b4SHong Zhang   const PetscInt  s = tab->s;
32463a6f1b4SHong Zhang   const PetscReal *A = tab->A,*c = tab->c;
32563a6f1b4SHong Zhang   PetscScalar     *w = rk->work;
32663a6f1b4SHong Zhang   PetscInt        i,j,k;
32763a6f1b4SHong Zhang   PetscReal       t = ts->ptime,h = ts->time_step;
32863a6f1b4SHong Zhang 
329362febeeSStefano Zampini   PetscFunctionBegin;
33063a6f1b4SHong Zhang   for (k=0; k<rk->dtratio; k++) {
331*9566063dSJacob Faibussowitsch     PetscCall(VecGetSubVector(ts->vec_sol,rk->is_fast,&Xfast));
33263a6f1b4SHong Zhang     for (i=0; i<s; i++) {
333*9566063dSJacob Faibussowitsch       PetscCall(VecGetSubVector(YdotRHS[i],rk->is_fast,&YdotRHS_fast[i]));
33463a6f1b4SHong Zhang     }
335a5b23f4aSJose E. Roman     /* propagate fast component using small time steps */
33663a6f1b4SHong Zhang     for (i=0; i<s; i++) {
33763a6f1b4SHong Zhang       /* stage value for slow components */
338*9566063dSJacob Faibussowitsch       PetscCall(TSInterpolate_RK_MultirateSplit(rk->ts_root,t+k*h/rk->dtratio+h/rk->dtratio*c[i],Y[i]));
33963a6f1b4SHong Zhang       currentlevelts = rk->ts_root;
34063a6f1b4SHong Zhang       while (currentlevelts != ts) { /* all the slow parts need to be interpolated separated */
34163a6f1b4SHong Zhang         currentlevelts = ((TS_RK*)currentlevelts->data)->subts_fast;
342*9566063dSJacob Faibussowitsch         PetscCall(TSInterpolate_RK_MultirateSplit(currentlevelts,t+k*h/rk->dtratio+h/rk->dtratio*c[i],Y[i]));
34363a6f1b4SHong Zhang       }
34463a6f1b4SHong Zhang       for (j=0; j<i; j++) w[j] = h/rk->dtratio*A[i*s+j];
34563a6f1b4SHong Zhang       subrk_fast->stage_time = t + h/rk->dtratio*c[i];
346*9566063dSJacob Faibussowitsch       PetscCall(TSPreStage(subts_fast,subrk_fast->stage_time));
34763a6f1b4SHong Zhang       /* stage value for fast components */
348*9566063dSJacob Faibussowitsch       PetscCall(VecGetSubVector(Y[i],rk->is_fast,&Yfast));
349*9566063dSJacob Faibussowitsch       PetscCall(VecCopy(Xfast,Yfast));
350*9566063dSJacob Faibussowitsch       PetscCall(VecMAXPY(Yfast,i,w,YdotRHS_fast));
351*9566063dSJacob Faibussowitsch       PetscCall(VecRestoreSubVector(Y[i],rk->is_fast,&Yfast));
352*9566063dSJacob Faibussowitsch       PetscCall(TSPostStage(subts_fast,subrk_fast->stage_time,i,Y));
35363a6f1b4SHong Zhang       /* compute the stage RHS for fast components */
354*9566063dSJacob Faibussowitsch       PetscCall(TSComputeRHSFunction(subts_fast,t+k*h*rk->dtratio+h/rk->dtratio*c[i],Y[i],YdotRHS_fast[i]));
35563a6f1b4SHong Zhang     }
356*9566063dSJacob Faibussowitsch     PetscCall(VecRestoreSubVector(ts->vec_sol,rk->is_fast,&Xfast));
35763a6f1b4SHong Zhang     /* update the value of fast components using fast time step */
35863a6f1b4SHong Zhang     rk->slow = PETSC_FALSE;
359*9566063dSJacob Faibussowitsch     PetscCall(TSEvaluateStep_RK_MultirateSplit(ts,tab->order,ts->vec_sol,NULL));
36063a6f1b4SHong Zhang     for (i=0; i<s; i++) {
361*9566063dSJacob Faibussowitsch       PetscCall(VecRestoreSubVector(YdotRHS[i],rk->is_fast,&YdotRHS_fast[i]));
36263a6f1b4SHong Zhang     }
36363a6f1b4SHong Zhang 
36463a6f1b4SHong Zhang     if (subrk_fast->subts_fast) {
36563a6f1b4SHong Zhang       subts_fast->ptime = t+k*h/rk->dtratio;
36663a6f1b4SHong Zhang       subts_fast->time_step = h/rk->dtratio;
367*9566063dSJacob Faibussowitsch       PetscCall(TSStepRefine_RK_MultirateSplit(subts_fast));
36863a6f1b4SHong Zhang     }
36963a6f1b4SHong Zhang     /* update the fast components of the solution */
370*9566063dSJacob Faibussowitsch     PetscCall(VecGetSubVector(ts->vec_sol,rk->is_fast,&Xfast));
371*9566063dSJacob Faibussowitsch     PetscCall(VecISCopy(rk->X0,rk->is_fast,SCATTER_FORWARD,Xfast));
372*9566063dSJacob Faibussowitsch     PetscCall(VecRestoreSubVector(ts->vec_sol,rk->is_fast,&Xfast));
37363a6f1b4SHong Zhang   }
37463a6f1b4SHong Zhang   PetscFunctionReturn(0);
37563a6f1b4SHong Zhang }
37663a6f1b4SHong Zhang 
37702555c94SHong Zhang static PetscErrorCode TSStep_RK_MultirateSplit(TS ts)
378474dd773SHong Zhang {
379474dd773SHong Zhang   TS_RK           *rk = (TS_RK*)ts->data;
380474dd773SHong Zhang   RKTableau       tab = rk->tableau;
381474dd773SHong Zhang   Vec             *Y = rk->Y,*YdotRHS = rk->YdotRHS;
38263a6f1b4SHong Zhang   Vec             *YdotRHS_fast = rk->YdotRHS_fast,*YdotRHS_slow = rk->YdotRHS_slow;
383474dd773SHong Zhang   Vec             Yslow,Yfast; /* subvectors store the stges of slow components and fast components respectively                           */
384474dd773SHong Zhang   const PetscInt  s = tab->s;
385474dd773SHong Zhang   const PetscReal *A = tab->A,*c = tab->c;
386474dd773SHong Zhang   PetscScalar     *w = rk->work;
38763a6f1b4SHong Zhang   PetscInt        i,j;
388474dd773SHong Zhang   PetscReal       next_time_step = ts->time_step,t = ts->ptime,h = ts->time_step;
389474dd773SHong Zhang 
390474dd773SHong Zhang   PetscFunctionBegin;
391474dd773SHong Zhang   rk->status = TS_STEP_INCOMPLETE;
392474dd773SHong Zhang   for (i=0; i<s; i++) {
393*9566063dSJacob Faibussowitsch     PetscCall(VecGetSubVector(YdotRHS[i],rk->is_slow,&YdotRHS_slow[i]));
394*9566063dSJacob Faibussowitsch     PetscCall(VecGetSubVector(YdotRHS[i],rk->is_fast,&YdotRHS_fast[i]));
395474dd773SHong Zhang   }
396*9566063dSJacob Faibussowitsch   PetscCall(VecCopy(ts->vec_sol,rk->X0));
397a5b23f4aSJose E. Roman   /* propagate both slow and fast components using large time steps */
398474dd773SHong Zhang   for (i=0; i<s; i++) {
399474dd773SHong Zhang     rk->stage_time = t + h*c[i];
400*9566063dSJacob Faibussowitsch     PetscCall(TSPreStage(ts,rk->stage_time));
401*9566063dSJacob Faibussowitsch     PetscCall(VecCopy(ts->vec_sol,Y[i]));
402*9566063dSJacob Faibussowitsch     PetscCall(VecGetSubVector(Y[i],rk->is_fast,&Yfast));
403*9566063dSJacob Faibussowitsch     PetscCall(VecGetSubVector(Y[i],rk->is_slow,&Yslow));
404474dd773SHong Zhang     for (j=0; j<i; j++) w[j] = h*A[i*s+j];
405*9566063dSJacob Faibussowitsch     PetscCall(VecMAXPY(Yfast,i,w,YdotRHS_fast));
406*9566063dSJacob Faibussowitsch     PetscCall(VecMAXPY(Yslow,i,w,YdotRHS_slow));
407*9566063dSJacob Faibussowitsch     PetscCall(VecRestoreSubVector(Y[i],rk->is_fast,&Yfast));
408*9566063dSJacob Faibussowitsch     PetscCall(VecRestoreSubVector(Y[i],rk->is_slow,&Yslow));
409*9566063dSJacob Faibussowitsch     PetscCall(TSPostStage(ts,rk->stage_time,i,Y));
410*9566063dSJacob Faibussowitsch     PetscCall(TSComputeRHSFunction(rk->subts_slow,t+h*c[i],Y[i],YdotRHS_slow[i]));
411*9566063dSJacob Faibussowitsch     PetscCall(TSComputeRHSFunction(rk->subts_fast,t+h*c[i],Y[i],YdotRHS_fast[i]));
412474dd773SHong Zhang   }
413474dd773SHong Zhang   rk->slow = PETSC_TRUE;
41463a6f1b4SHong Zhang   /* update the slow components of the solution using slow time step */
415*9566063dSJacob Faibussowitsch   PetscCall(TSEvaluateStep_RK_MultirateSplit(ts,tab->order,ts->vec_sol,NULL));
41663a6f1b4SHong Zhang   /* YdotRHS will be used for interpolation during refinement */
417474dd773SHong Zhang   for (i=0; i<s; i++) {
418*9566063dSJacob Faibussowitsch     PetscCall(VecRestoreSubVector(YdotRHS[i],rk->is_slow,&YdotRHS_slow[i]));
419*9566063dSJacob Faibussowitsch     PetscCall(VecRestoreSubVector(YdotRHS[i],rk->is_fast,&YdotRHS_fast[i]));
420474dd773SHong Zhang   }
42163a6f1b4SHong Zhang 
422*9566063dSJacob Faibussowitsch   PetscCall(TSStepRefine_RK_MultirateSplit(ts));
42363a6f1b4SHong Zhang 
424bb42530cSHong Zhang   ts->ptime = t + ts->time_step;
425474dd773SHong Zhang   ts->time_step = next_time_step;
426474dd773SHong Zhang   rk->status = TS_STEP_COMPLETE;
427474dd773SHong Zhang   PetscFunctionReturn(0);
428474dd773SHong Zhang }
429474dd773SHong Zhang 
43002555c94SHong Zhang static PetscErrorCode TSSetUp_RK_MultirateSplit(TS ts)
4310fe4d17eSHong Zhang {
4320fe4d17eSHong Zhang   TS_RK          *rk = (TS_RK*)ts->data,*nextlevelrk,*currentlevelrk;
4330fe4d17eSHong Zhang   TS             nextlevelts;
4340fe4d17eSHong Zhang   Vec            X0;
4350fe4d17eSHong Zhang 
4360fe4d17eSHong Zhang   PetscFunctionBegin;
437*9566063dSJacob Faibussowitsch   PetscCall(TSRHSSplitGetIS(ts,"slow",&rk->is_slow));
438*9566063dSJacob Faibussowitsch   PetscCall(TSRHSSplitGetIS(ts,"fast",&rk->is_fast));
4393c633725SBarry Smith   PetscCheck(rk->is_slow && rk->is_fast,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");
440*9566063dSJacob Faibussowitsch   PetscCall(TSRHSSplitGetSubTS(ts,"slow",&rk->subts_slow));
441*9566063dSJacob Faibussowitsch   PetscCall(TSRHSSplitGetSubTS(ts,"fast",&rk->subts_fast));
4423c633725SBarry Smith   PetscCheck(rk->subts_slow && rk->subts_fast,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");
4430fe4d17eSHong Zhang 
444*9566063dSJacob Faibussowitsch   PetscCall(VecDuplicate(ts->vec_sol,&X0));
4450fe4d17eSHong Zhang   /* The TS at each level share the same tableau, work array, solution vector, stage values and stage derivatives */
4460fe4d17eSHong Zhang   currentlevelrk = rk;
4470fe4d17eSHong Zhang   while (currentlevelrk->subts_fast) {
448*9566063dSJacob Faibussowitsch     PetscCall(PetscMalloc1(rk->tableau->s,&currentlevelrk->YdotRHS_fast));
449*9566063dSJacob Faibussowitsch     PetscCall(PetscMalloc1(rk->tableau->s,&currentlevelrk->YdotRHS_slow));
450*9566063dSJacob Faibussowitsch     PetscCall(PetscObjectReference((PetscObject)X0));
4510fe4d17eSHong Zhang     currentlevelrk->X0 = X0;
4520fe4d17eSHong Zhang     currentlevelrk->ts_root = ts;
4530fe4d17eSHong Zhang 
4540fe4d17eSHong Zhang     /* set up the ts for the slow part */
4550fe4d17eSHong Zhang     nextlevelts = currentlevelrk->subts_slow;
456*9566063dSJacob Faibussowitsch     PetscCall(PetscNewLog(nextlevelts,&nextlevelrk));
4570fe4d17eSHong Zhang     nextlevelrk->tableau = rk->tableau;
4580fe4d17eSHong Zhang     nextlevelrk->work = rk->work;
4590fe4d17eSHong Zhang     nextlevelrk->Y = rk->Y;
4600fe4d17eSHong Zhang     nextlevelrk->YdotRHS = rk->YdotRHS;
4610fe4d17eSHong Zhang     nextlevelts->data = (void*)nextlevelrk;
462*9566063dSJacob Faibussowitsch     PetscCall(TSCopyDM(ts,nextlevelts));
463*9566063dSJacob Faibussowitsch     PetscCall(TSSetSolution(nextlevelts,ts->vec_sol));
4640fe4d17eSHong Zhang 
4650fe4d17eSHong Zhang     /* set up the ts for the fast part */
4660fe4d17eSHong Zhang     nextlevelts = currentlevelrk->subts_fast;
467*9566063dSJacob Faibussowitsch     PetscCall(PetscNewLog(nextlevelts,&nextlevelrk));
4680fe4d17eSHong Zhang     nextlevelrk->tableau = rk->tableau;
4690fe4d17eSHong Zhang     nextlevelrk->work = rk->work;
4700fe4d17eSHong Zhang     nextlevelrk->Y = rk->Y;
4710fe4d17eSHong Zhang     nextlevelrk->YdotRHS = rk->YdotRHS;
4720fe4d17eSHong Zhang     nextlevelrk->dtratio = rk->dtratio;
473*9566063dSJacob Faibussowitsch     PetscCall(TSRHSSplitGetIS(nextlevelts,"slow",&nextlevelrk->is_slow));
474*9566063dSJacob Faibussowitsch     PetscCall(TSRHSSplitGetSubTS(nextlevelts,"slow",&nextlevelrk->subts_slow));
475*9566063dSJacob Faibussowitsch     PetscCall(TSRHSSplitGetIS(nextlevelts,"fast",&nextlevelrk->is_fast));
476*9566063dSJacob Faibussowitsch     PetscCall(TSRHSSplitGetSubTS(nextlevelts,"fast",&nextlevelrk->subts_fast));
4770fe4d17eSHong Zhang     nextlevelts->data = (void*)nextlevelrk;
478*9566063dSJacob Faibussowitsch     PetscCall(TSCopyDM(ts,nextlevelts));
479*9566063dSJacob Faibussowitsch     PetscCall(TSSetSolution(nextlevelts,ts->vec_sol));
4800fe4d17eSHong Zhang 
4810fe4d17eSHong Zhang     currentlevelrk = nextlevelrk;
4820fe4d17eSHong Zhang   }
483*9566063dSJacob Faibussowitsch   PetscCall(VecDestroy(&X0));
4840fe4d17eSHong Zhang 
48502555c94SHong Zhang   ts->ops->step         = TSStep_RK_MultirateSplit;
48602555c94SHong Zhang   ts->ops->evaluatestep = TSEvaluateStep_RK_MultirateSplit;
48702555c94SHong Zhang   ts->ops->interpolate  = TSInterpolate_RK_MultirateSplit;
4880fe4d17eSHong Zhang   PetscFunctionReturn(0);
4890fe4d17eSHong Zhang }
4900fe4d17eSHong Zhang 
49121052c0cSHong Zhang PetscErrorCode TSRKSetMultirate_RK(TS ts,PetscBool use_multirate)
492474dd773SHong Zhang {
493474dd773SHong Zhang   TS_RK          *rk = (TS_RK*)ts->data;
494474dd773SHong Zhang 
495474dd773SHong Zhang   PetscFunctionBegin;
496474dd773SHong Zhang   PetscValidHeaderSpecific(ts,TS_CLASSID,1);
4970fe4d17eSHong Zhang   rk->use_multirate = use_multirate;
4980fe4d17eSHong Zhang   if (use_multirate) {
499474dd773SHong Zhang     rk->dtratio = 2;
500*9566063dSJacob Faibussowitsch     PetscCall(PetscObjectComposeFunction((PetscObject)ts,"TSSetUp_RK_MultirateSplit_C",TSSetUp_RK_MultirateSplit));
501*9566063dSJacob Faibussowitsch     PetscCall(PetscObjectComposeFunction((PetscObject)ts,"TSReset_RK_MultirateSplit_C",TSReset_RK_MultirateSplit));
502*9566063dSJacob Faibussowitsch     PetscCall(PetscObjectComposeFunction((PetscObject)ts,"TSSetUp_RK_MultirateNonsplit_C",TSSetUp_RK_MultirateNonsplit));
503*9566063dSJacob Faibussowitsch     PetscCall(PetscObjectComposeFunction((PetscObject)ts,"TSReset_RK_MultirateNonsplit_C",TSReset_RK_MultirateNonsplit));
5040fe4d17eSHong Zhang   } else {
5050fe4d17eSHong Zhang     rk->dtratio = 0;
506*9566063dSJacob Faibussowitsch     PetscCall(PetscObjectComposeFunction((PetscObject)ts,"TSSetUp_RK_MultirateSplit_C",NULL));
507*9566063dSJacob Faibussowitsch     PetscCall(PetscObjectComposeFunction((PetscObject)ts,"TSReset_RK_MultirateSplit_C",NULL));
508*9566063dSJacob Faibussowitsch     PetscCall(PetscObjectComposeFunction((PetscObject)ts,"TSSetUp_RK_MultirateNonsplit_C",NULL));
509*9566063dSJacob Faibussowitsch     PetscCall(PetscObjectComposeFunction((PetscObject)ts,"TSReset_RK_MultirateNonsplit_C",NULL));
510474dd773SHong Zhang   }
511474dd773SHong Zhang   PetscFunctionReturn(0);
512474dd773SHong Zhang }
513474dd773SHong Zhang 
51421052c0cSHong Zhang PetscErrorCode TSRKGetMultirate_RK(TS ts,PetscBool *use_multirate)
5150fe4d17eSHong Zhang {
5160fe4d17eSHong Zhang   TS_RK *rk = (TS_RK*)ts->data;
5170fe4d17eSHong Zhang 
5180fe4d17eSHong Zhang   PetscFunctionBegin;
5190fe4d17eSHong Zhang   PetscValidHeaderSpecific(ts,TS_CLASSID,1);
5200fe4d17eSHong Zhang   *use_multirate = rk->use_multirate;
5210fe4d17eSHong Zhang   PetscFunctionReturn(0);
5220fe4d17eSHong Zhang }
523