xref: /petsc/src/ts/impls/explicit/rk/mrk.c (revision bb42530c4a0604e316f9f88124d6eb2dd52b7c32)
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>
17474dd773SHong Zhang 
18474dd773SHong Zhang static TSRKType TSMRKDefault = TSRK2A;
19474dd773SHong Zhang 
20474dd773SHong Zhang static PetscErrorCode TSSetUp_MRKSPLIT(TS ts)
21474dd773SHong Zhang {
22*bb42530cSHong Zhang   TS_RK          *rk = (TS_RK*)ts->data,*nextlevelrk,*currentlevelrk;
23*bb42530cSHong Zhang   TS             nextlevelts;
24474dd773SHong Zhang   DM             dm,subdm,newdm;
25474dd773SHong Zhang   PetscErrorCode ierr;
26474dd773SHong Zhang 
27474dd773SHong Zhang   PetscFunctionBegin;
28474dd773SHong Zhang   ierr = TSRHSSplitGetIS(ts,"slow",&rk->is_slow);CHKERRQ(ierr);
29474dd773SHong Zhang   ierr = TSRHSSplitGetIS(ts,"fast",&rk->is_fast);CHKERRQ(ierr);
30474dd773SHong Zhang   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");
31474dd773SHong Zhang   ierr = TSRHSSplitGetSubTS(ts,"slow",&rk->subts_slow);CHKERRQ(ierr);
32474dd773SHong Zhang   ierr = TSRHSSplitGetSubTS(ts,"fast",&rk->subts_fast);CHKERRQ(ierr);
33474dd773SHong Zhang   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");
34474dd773SHong Zhang 
35474dd773SHong Zhang   /* Only copy */
36474dd773SHong Zhang   ierr = TSGetDM(ts,&dm);CHKERRQ(ierr);
37474dd773SHong Zhang   ierr = DMClone(dm,&newdm);CHKERRQ(ierr);
38474dd773SHong Zhang   ierr = TSGetDM(rk->subts_fast,&subdm);CHKERRQ(ierr);
39474dd773SHong Zhang   ierr = DMCopyDMTS(subdm,newdm);CHKERRQ(ierr);
40474dd773SHong Zhang   ierr = DMCopyDMSNES(subdm,newdm);CHKERRQ(ierr);
41474dd773SHong Zhang   ierr = TSSetDM(rk->subts_fast,newdm);CHKERRQ(ierr);
42474dd773SHong Zhang   ierr = DMDestroy(&newdm);CHKERRQ(ierr);
43474dd773SHong Zhang   ierr = DMClone(dm,&newdm);CHKERRQ(ierr);
44474dd773SHong Zhang   ierr = TSGetDM(rk->subts_slow,&subdm);CHKERRQ(ierr);
45474dd773SHong Zhang   ierr = DMCopyDMTS(subdm,newdm);CHKERRQ(ierr);
46474dd773SHong Zhang   ierr = DMCopyDMSNES(subdm,newdm);CHKERRQ(ierr);
47474dd773SHong Zhang   ierr = TSSetDM(rk->subts_slow,newdm);CHKERRQ(ierr);
48474dd773SHong Zhang   ierr = DMDestroy(&newdm);CHKERRQ(ierr);
49474dd773SHong Zhang 
50474dd773SHong Zhang   ierr = VecDuplicate(ts->vec_sol,&rk->X0);CHKERRQ(ierr);
51*bb42530cSHong Zhang   /* The TS at each level share the same tableau, work array, solution vector, stage values and stage derivatives */
52*bb42530cSHong Zhang   currentlevelrk = rk;
53*bb42530cSHong Zhang   while (currentlevelrk->subts_fast) {
54*bb42530cSHong Zhang     /* set up the ts for the slow part */
55*bb42530cSHong Zhang     nextlevelts = currentlevelrk->subts_slow;
56*bb42530cSHong Zhang     ierr = PetscNewLog(nextlevelts,&nextlevelrk);CHKERRQ(ierr);
57*bb42530cSHong Zhang     nextlevelrk->tableau = rk->tableau;
58*bb42530cSHong Zhang     nextlevelrk->work = rk->work;
59*bb42530cSHong Zhang     nextlevelrk->Y = rk->Y;
60*bb42530cSHong Zhang     nextlevelrk->X0 = rk->X0;
61*bb42530cSHong Zhang     ierr = PetscMalloc1(rk->tableau->s,&nextlevelrk->YdotRHS);CHKERRQ(ierr);
62*bb42530cSHong Zhang     nextlevelts->data = (void*)nextlevelrk;
63*bb42530cSHong Zhang     ierr = TSSetSolution(nextlevelts,ts->vec_sol);CHKERRQ(ierr);
64*bb42530cSHong Zhang 
65*bb42530cSHong Zhang     /* set up the ts for the fast part */
66*bb42530cSHong Zhang     nextlevelts = currentlevelrk->subts_fast;
67*bb42530cSHong Zhang     ierr = PetscNewLog(nextlevelts,&nextlevelrk);CHKERRQ(ierr);
68*bb42530cSHong Zhang     nextlevelrk->tableau = rk->tableau;
69*bb42530cSHong Zhang     nextlevelrk->work = rk->work;
70*bb42530cSHong Zhang     nextlevelrk->Y = rk->Y;
71*bb42530cSHong Zhang     nextlevelrk->X0 = rk->X0;
72*bb42530cSHong Zhang     nextlevelrk->dtratio = rk->dtratio;
73*bb42530cSHong Zhang     ierr = PetscMalloc1(rk->tableau->s,&nextlevelrk->YdotRHS);CHKERRQ(ierr);
74*bb42530cSHong Zhang     ierr = TSRHSSplitGetIS(nextlevelts,"slow",&nextlevelrk->is_slow);CHKERRQ(ierr);
75*bb42530cSHong Zhang     ierr = TSRHSSplitGetSubTS(nextlevelts,"slow",&nextlevelrk->subts_slow);CHKERRQ(ierr);
76*bb42530cSHong Zhang     ierr = TSRHSSplitGetIS(nextlevelts,"fast",&nextlevelrk->is_fast);CHKERRQ(ierr);
77*bb42530cSHong Zhang     ierr = TSRHSSplitGetSubTS(nextlevelts,"fast",&nextlevelrk->subts_fast);CHKERRQ(ierr);
78*bb42530cSHong Zhang     nextlevelts->data = (void*)nextlevelrk;
79*bb42530cSHong Zhang     ierr = TSSetSolution(nextlevelts,ts->vec_sol);CHKERRQ(ierr);
80*bb42530cSHong Zhang 
81*bb42530cSHong Zhang     currentlevelrk = nextlevelrk;
82*bb42530cSHong Zhang   }
83474dd773SHong Zhang   PetscFunctionReturn(0);
84474dd773SHong Zhang }
85474dd773SHong Zhang 
86474dd773SHong Zhang static PetscErrorCode TSReset_MRKSPLIT(TS ts)
87474dd773SHong Zhang {
88474dd773SHong Zhang   TS_RK          *rk = (TS_RK*)ts->data;
89474dd773SHong Zhang   PetscErrorCode ierr;
90474dd773SHong Zhang 
91474dd773SHong Zhang   PetscFunctionBegin;
92*bb42530cSHong Zhang   if (rk->subts_slow) {
93*bb42530cSHong Zhang     ierr = TSReset_MRKSPLIT(rk->subts_slow);CHKERRQ(ierr);
94*bb42530cSHong Zhang     ierr = PetscFree(rk->subts_slow->data);CHKERRQ(ierr);
95*bb42530cSHong Zhang     rk->subts_slow = NULL;
96*bb42530cSHong Zhang   }
97*bb42530cSHong Zhang   if (rk->subts_fast) {
98*bb42530cSHong Zhang     ierr = VecDestroy(&rk->X0);CHKERRQ(ierr);
99*bb42530cSHong Zhang     ierr = TSReset_MRKSPLIT(rk->subts_fast);
100*bb42530cSHong Zhang     ierr = PetscFree(rk->subts_fast->data);CHKERRQ(ierr);
101*bb42530cSHong Zhang     rk->subts_fast = NULL;
102*bb42530cSHong Zhang   }
103474dd773SHong Zhang   ierr = PetscFree(rk->YdotRHS_fast);CHKERRQ(ierr);
104474dd773SHong Zhang   ierr = PetscFree(rk->YdotRHS_slow);CHKERRQ(ierr);
105*bb42530cSHong Zhang   ierr = PetscFree(rk->YdotRHS);CHKERRQ(ierr);
106474dd773SHong Zhang   PetscFunctionReturn(0);
107474dd773SHong Zhang }
108474dd773SHong Zhang 
109474dd773SHong Zhang static PetscErrorCode TSSetUp_MRKNONSPLIT(TS ts)
110474dd773SHong Zhang {
111474dd773SHong Zhang   TS_RK          *rk = (TS_RK*)ts->data;
112474dd773SHong Zhang   RKTableau       tab  = rk->tableau;
113474dd773SHong Zhang   PetscErrorCode ierr;
114474dd773SHong Zhang 
115474dd773SHong Zhang   PetscFunctionBegin;
116474dd773SHong Zhang   ierr = TSRHSSplitGetIS(ts,"slow",&rk->is_slow);CHKERRQ(ierr);
117474dd773SHong Zhang   ierr = TSRHSSplitGetIS(ts,"fast",&rk->is_fast);CHKERRQ(ierr);
118474dd773SHong Zhang   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");
119474dd773SHong Zhang   ierr = VecDuplicate(ts->vec_sol,&rk->X0);CHKERRQ(ierr);
120474dd773SHong Zhang   ierr = VecDuplicateVecs(ts->vec_sol,tab->s,&rk->YdotRHS_slow);CHKERRQ(ierr);
121474dd773SHong Zhang   PetscFunctionReturn(0);
122474dd773SHong Zhang }
123474dd773SHong Zhang 
124474dd773SHong Zhang static PetscErrorCode TSReset_MRKNONSPLIT(TS ts)
125474dd773SHong Zhang {
126474dd773SHong Zhang   TS_RK          *rk = (TS_RK*)ts->data;
127474dd773SHong Zhang   RKTableau      tab  = rk->tableau;
128474dd773SHong Zhang   PetscErrorCode ierr;
129474dd773SHong Zhang 
130474dd773SHong Zhang   PetscFunctionBegin;
131474dd773SHong Zhang   ierr = VecDestroy(&rk->X0);CHKERRQ(ierr);
132474dd773SHong Zhang   ierr = VecDestroyVecs(tab->s,&rk->YdotRHS_slow);CHKERRQ(ierr);
133474dd773SHong Zhang   PetscFunctionReturn(0);
134474dd773SHong Zhang }
135474dd773SHong Zhang 
136474dd773SHong Zhang static PetscErrorCode TSInterpolate_MRKNONSPLIT(TS ts,PetscReal itime,Vec X)
137474dd773SHong Zhang {
138474dd773SHong Zhang   TS_RK            *rk = (TS_RK*)ts->data;
139474dd773SHong Zhang   PetscInt         s  = rk->tableau->s,p = rk->tableau->p,i,j;
140474dd773SHong Zhang   PetscReal        h;
141474dd773SHong Zhang   PetscReal        tt,t;
142474dd773SHong Zhang   PetscScalar      *b;
143474dd773SHong Zhang   const PetscReal  *B = rk->tableau->binterp;
144474dd773SHong Zhang   PetscErrorCode   ierr;
145474dd773SHong Zhang 
146474dd773SHong Zhang   PetscFunctionBegin;
147474dd773SHong Zhang   if (!B) SETERRQ1(PetscObjectComm((PetscObject)ts),PETSC_ERR_SUP,"TSRK %s does not have an interpolation formula",rk->tableau->name);
148474dd773SHong Zhang 
149474dd773SHong Zhang   switch (rk->status) {
150474dd773SHong Zhang     case TS_STEP_INCOMPLETE:
151474dd773SHong Zhang     case TS_STEP_PENDING:
152474dd773SHong Zhang       h = ts->time_step;
153474dd773SHong Zhang       t = (itime - ts->ptime)/h;
154474dd773SHong Zhang       break;
155474dd773SHong Zhang     case TS_STEP_COMPLETE:
156474dd773SHong Zhang       h = ts->ptime - ts->ptime_prev;
157474dd773SHong Zhang       t = (itime - ts->ptime)/h + 1; /* In the interval [0,1] */
158474dd773SHong Zhang       break;
159474dd773SHong Zhang     default: SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_PLIB,"Invalid TSStepStatus");
160474dd773SHong Zhang   }
161474dd773SHong Zhang   ierr = PetscMalloc1(s,&b);CHKERRQ(ierr);
162474dd773SHong Zhang   for (i=0; i<s; i++) b[i] = 0;
163474dd773SHong Zhang   for (j=0,tt=t; j<p; j++,tt*=t) {
164474dd773SHong Zhang     for (i=0; i<s; i++) {
165474dd773SHong Zhang       b[i]  += h * B[i*p+j] * tt;
166474dd773SHong Zhang     }
167474dd773SHong Zhang   }
168474dd773SHong Zhang   ierr = VecCopy(rk->X0,X);CHKERRQ(ierr);
169474dd773SHong Zhang   ierr = VecMAXPY(X,s,b,rk->YdotRHS_slow);CHKERRQ(ierr);
170474dd773SHong Zhang   ierr = PetscFree(b);CHKERRQ(ierr);
171474dd773SHong Zhang   PetscFunctionReturn(0);
172474dd773SHong Zhang }
173474dd773SHong Zhang 
174474dd773SHong Zhang static PetscErrorCode TSInterpolate_MRKSPLIT(TS ts,PetscReal itime,Vec X)
175474dd773SHong Zhang {
176474dd773SHong Zhang   TS_RK            *rk = (TS_RK*)ts->data;
177474dd773SHong Zhang   PetscInt         s  = rk->tableau->s,p = rk->tableau->p,i,j;
178474dd773SHong Zhang   Vec              Yslow;    /* vector holds the slow components in Y[0] */
179474dd773SHong Zhang   PetscReal        h;
180474dd773SHong Zhang   PetscReal        tt,t;
181474dd773SHong Zhang   PetscScalar      *b;
182474dd773SHong Zhang   const PetscReal  *B = rk->tableau->binterp;
183474dd773SHong Zhang   PetscErrorCode   ierr;
184474dd773SHong Zhang 
185474dd773SHong Zhang   PetscFunctionBegin;
186474dd773SHong Zhang   if (!B) SETERRQ1(PetscObjectComm((PetscObject)ts),PETSC_ERR_SUP,"TSRK %s does not have an interpolation formula",rk->tableau->name);
187474dd773SHong Zhang 
188474dd773SHong Zhang   switch (rk->status) {
189474dd773SHong Zhang     case TS_STEP_INCOMPLETE:
190474dd773SHong Zhang     case TS_STEP_PENDING:
191474dd773SHong Zhang       h = ts->time_step;
192474dd773SHong Zhang       t = (itime - ts->ptime)/h;
193474dd773SHong Zhang       break;
194474dd773SHong Zhang     case TS_STEP_COMPLETE:
195474dd773SHong Zhang       h = ts->ptime - ts->ptime_prev;
196474dd773SHong Zhang       t = (itime - ts->ptime)/h + 1; /* In the interval [0,1] */
197474dd773SHong Zhang       break;
198474dd773SHong Zhang     default: SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_PLIB,"Invalid TSStepStatus");
199474dd773SHong Zhang   }
200474dd773SHong Zhang   ierr = PetscMalloc1(s,&b);CHKERRQ(ierr);
201474dd773SHong Zhang   for (i=0; i<s; i++) b[i] = 0;
202474dd773SHong Zhang   for (j=0,tt=t; j<p; j++,tt*=t) {
203474dd773SHong Zhang     for (i=0; i<s; i++) {
204474dd773SHong Zhang       b[i]  += h * B[i*p+j] * tt;
205474dd773SHong Zhang     }
206474dd773SHong Zhang   }
207474dd773SHong Zhang   ierr = VecGetSubVector(rk->X0,rk->is_slow,&Yslow);CHKERRQ(ierr);
208474dd773SHong Zhang   ierr = VecCopy(Yslow,X);CHKERRQ(ierr);
209*bb42530cSHong Zhang   ierr = VecMAXPY(X,s,b,((TS_RK*)rk->subts_slow->data)->YdotRHS);CHKERRQ(ierr);
210474dd773SHong Zhang   ierr = VecRestoreSubVector(rk->X0,rk->is_slow,&Yslow);CHKERRQ(ierr);
211474dd773SHong Zhang   ierr = PetscFree(b);CHKERRQ(ierr);
212474dd773SHong Zhang   PetscFunctionReturn(0);
213474dd773SHong Zhang }
214474dd773SHong Zhang 
215474dd773SHong Zhang static PetscErrorCode TSStep_MRKNONSPLIT(TS ts)
216474dd773SHong Zhang {
217474dd773SHong Zhang   TS_RK             *rk = (TS_RK*)ts->data;
218474dd773SHong Zhang   RKTableau         tab  = rk->tableau;
219474dd773SHong Zhang   Vec               *Y = rk->Y,*YdotRHS = rk->YdotRHS,*YdotRHS_slow = rk->YdotRHS_slow;
220474dd773SHong Zhang   Vec               stage_slow,sol_slow;   /* vectors store the slow components */
221474dd773SHong Zhang   Vec               subvec_slow;           /* sub vector to store the slow components */
222474dd773SHong Zhang   const PetscInt    s = tab->s;
223474dd773SHong Zhang   const PetscReal   *A = tab->A,*c = tab->c;
224474dd773SHong Zhang   PetscScalar       *w = rk->work;
225474dd773SHong Zhang   PetscInt          i,j,k,dtratio = rk->dtratio;
226474dd773SHong Zhang   PetscReal         next_time_step = ts->time_step,t = ts->ptime,h = ts->time_step;
227474dd773SHong Zhang   PetscErrorCode    ierr;
228474dd773SHong Zhang 
229474dd773SHong Zhang   PetscFunctionBegin;
230474dd773SHong Zhang   rk->status = TS_STEP_INCOMPLETE;
231474dd773SHong Zhang   ierr = VecDuplicate(ts->vec_sol,&stage_slow);CHKERRQ(ierr);
232474dd773SHong Zhang   ierr = VecDuplicate(ts->vec_sol,&sol_slow);CHKERRQ(ierr);
233474dd773SHong Zhang   ierr = VecCopy(ts->vec_sol,rk->X0);CHKERRQ(ierr);
234474dd773SHong Zhang   for (i=0; i<s; i++) {
235474dd773SHong Zhang     rk->stage_time = t + h*c[i];
236474dd773SHong Zhang     ierr = TSPreStage(ts,rk->stage_time);CHKERRQ(ierr);
237474dd773SHong Zhang     ierr = VecCopy(ts->vec_sol,Y[i]);CHKERRQ(ierr);
238474dd773SHong Zhang     for (j=0; j<i; j++) w[j] = h*A[i*s+j];
239474dd773SHong Zhang     ierr = VecMAXPY(Y[i],i,w,YdotRHS_slow);CHKERRQ(ierr);
240474dd773SHong Zhang     ierr = TSPostStage(ts,rk->stage_time,i,Y); CHKERRQ(ierr);
241474dd773SHong Zhang     /* compute the stage RHS */
242474dd773SHong Zhang     ierr = TSComputeRHSFunction(ts,t+h*c[i],Y[i],YdotRHS_slow[i]);CHKERRQ(ierr);
243474dd773SHong Zhang   }
244474dd773SHong Zhang   /* update the slow components in the solution */
245474dd773SHong Zhang   rk->YdotRHS = YdotRHS_slow;
246474dd773SHong Zhang   rk->dtratio = 1;
247474dd773SHong Zhang   ierr = TSEvaluateStep(ts,tab->order,sol_slow,NULL);CHKERRQ(ierr);
248474dd773SHong Zhang   rk->dtratio = dtratio;
249474dd773SHong Zhang   rk->YdotRHS = YdotRHS;
250474dd773SHong Zhang   for (k=0; k<rk->dtratio; k++) {
251474dd773SHong Zhang     for (i=0; i<s; i++) {
252474dd773SHong Zhang       rk->stage_time = t + h/rk->dtratio*c[i];
253474dd773SHong Zhang       ierr = TSPreStage(ts,rk->stage_time);CHKERRQ(ierr);
254474dd773SHong Zhang       /* 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 */
255474dd773SHong Zhang       ierr = VecCopy(ts->vec_sol,Y[i]);CHKERRQ(ierr);
256474dd773SHong Zhang       for (j=0; j<i; j++) w[j] = h/rk->dtratio*A[i*s+j];
257474dd773SHong Zhang       ierr = VecMAXPY(Y[i],i,w,YdotRHS);CHKERRQ(ierr);
258474dd773SHong Zhang       ierr = TSInterpolate_MRKNONSPLIT(ts,t+k*h/rk->dtratio+h/rk->dtratio*c[i],stage_slow);CHKERRQ(ierr);
259474dd773SHong Zhang       /* update the slow components in the stage value */
260474dd773SHong Zhang       ierr = VecGetSubVector(stage_slow,rk->is_slow,&subvec_slow);CHKERRQ(ierr);
261474dd773SHong Zhang       ierr = VecISCopy(Y[i],rk->is_slow,SCATTER_FORWARD,subvec_slow);CHKERRQ(ierr);
262474dd773SHong Zhang       ierr = VecRestoreSubVector(stage_slow,rk->is_slow,&subvec_slow);CHKERRQ(ierr);
263474dd773SHong Zhang       ierr = TSPostStage(ts,rk->stage_time,i,Y);CHKERRQ(ierr);
264474dd773SHong Zhang       /* compute the stage RHS */
265474dd773SHong Zhang       ierr = TSComputeRHSFunction(ts,t+k*h/rk->dtratio+h/rk->dtratio*c[i],Y[i],YdotRHS[i]);CHKERRQ(ierr);
266474dd773SHong Zhang     }
267474dd773SHong Zhang     ierr = TSEvaluateStep(ts,tab->order,ts->vec_sol,NULL);CHKERRQ(ierr);
268474dd773SHong Zhang   }
269474dd773SHong Zhang   /* update the slow components in the solution */
270474dd773SHong Zhang   ierr = VecGetSubVector(sol_slow,rk->is_slow,&subvec_slow);CHKERRQ(ierr);
271474dd773SHong Zhang   ierr = VecISCopy(ts->vec_sol,rk->is_slow,SCATTER_FORWARD,subvec_slow);CHKERRQ(ierr);
272474dd773SHong Zhang   ierr = VecRestoreSubVector(sol_slow,rk->is_slow,&subvec_slow);CHKERRQ(ierr);
273474dd773SHong Zhang 
274474dd773SHong Zhang   ts->ptime += ts->time_step;
275474dd773SHong Zhang   ts->time_step = next_time_step;
276474dd773SHong Zhang   rk->status = TS_STEP_COMPLETE;
277474dd773SHong Zhang   /* free memory of work vectors */
278474dd773SHong Zhang   ierr = VecDestroy(&stage_slow);CHKERRQ(ierr);
279474dd773SHong Zhang   ierr = VecDestroy(&sol_slow);CHKERRQ(ierr);
280474dd773SHong Zhang   PetscFunctionReturn(0);
281474dd773SHong Zhang }
282474dd773SHong Zhang 
283474dd773SHong Zhang /*
284474dd773SHong Zhang  This is for partitioned RHS multirate RK method
285474dd773SHong Zhang  The step completion formula is
286474dd773SHong Zhang 
287474dd773SHong Zhang  x1 = x0 + h b^T YdotRHS
288474dd773SHong Zhang 
289474dd773SHong Zhang */
290474dd773SHong Zhang static PetscErrorCode TSEvaluateStep_MRKSPLIT(TS ts,PetscInt order,Vec X,PetscBool *done)
291474dd773SHong Zhang {
292474dd773SHong Zhang   TS_RK           *rk = (TS_RK*)ts->data;
293474dd773SHong Zhang   RKTableau       tab  = rk->tableau;
294474dd773SHong Zhang   Vec             Xslow,Xfast;                  /* subvectors of X which store slow components and fast components respectively */
295474dd773SHong Zhang   PetscScalar     *w = rk->work;
296474dd773SHong Zhang   PetscReal       h = ts->time_step;
297474dd773SHong Zhang   PetscInt        s = tab->s,j;
298474dd773SHong Zhang   PetscErrorCode  ierr;
299474dd773SHong Zhang 
300474dd773SHong Zhang   PetscFunctionBegin;
301474dd773SHong Zhang   ierr = VecCopy(ts->vec_sol,X);CHKERRQ(ierr);
302474dd773SHong Zhang   if (rk->slow) {
303474dd773SHong Zhang     for (j=0; j<s; j++) w[j] = h*tab->b[j];
304474dd773SHong Zhang     ierr = VecGetSubVector(ts->vec_sol,rk->is_slow,&Xslow);CHKERRQ(ierr);
305*bb42530cSHong Zhang     ierr = VecMAXPY(Xslow,s,w,((TS_RK*)rk->subts_slow->data)->YdotRHS);CHKERRQ(ierr);
306474dd773SHong Zhang     ierr = VecRestoreSubVector(ts->vec_sol,rk->is_slow,&Xslow);CHKERRQ(ierr);;
307474dd773SHong Zhang   } else {
308474dd773SHong Zhang     for (j=0; j<s; j++) w[j] = h/rk->dtratio*tab->b[j];
309474dd773SHong Zhang     ierr = VecGetSubVector(X,rk->is_fast,&Xfast);CHKERRQ(ierr);
310*bb42530cSHong Zhang     ierr = VecMAXPY(Xfast,s,w,((TS_RK*)rk->subts_fast->data)->YdotRHS);CHKERRQ(ierr);
311474dd773SHong Zhang     ierr = VecRestoreSubVector(X,rk->is_fast,&Xfast);CHKERRQ(ierr);
312474dd773SHong Zhang   }
313474dd773SHong Zhang   PetscFunctionReturn(0);
314474dd773SHong Zhang }
315474dd773SHong Zhang 
316474dd773SHong Zhang static PetscErrorCode TSStep_MRKSPLIT(TS ts)
317474dd773SHong Zhang {
318474dd773SHong Zhang   TS_RK             *rk = (TS_RK*)ts->data;
319474dd773SHong Zhang   RKTableau         tab = rk->tableau;
320474dd773SHong Zhang   Vec               *Y = rk->Y,*YdotRHS = rk->YdotRHS;
321*bb42530cSHong Zhang   Vec               *YdotRHS_fast,*YdotRHS_slow;
322474dd773SHong Zhang   Vec               Yslow,Yfast;                         /* subvectors store the stges of slow components and fast components respectively                           */
323474dd773SHong Zhang   const PetscInt    s = tab->s;
324474dd773SHong Zhang   const PetscReal   *A = tab->A,*c = tab->c;
325474dd773SHong Zhang   PetscScalar       *w = rk->work;
326474dd773SHong Zhang   PetscInt          i,j,k;
327474dd773SHong Zhang   PetscReal         next_time_step = ts->time_step,t = ts->ptime,h = ts->time_step;
328474dd773SHong Zhang   PetscErrorCode    ierr;
329474dd773SHong Zhang 
330474dd773SHong Zhang   PetscFunctionBegin;
331474dd773SHong Zhang   rk->status = TS_STEP_INCOMPLETE;
332*bb42530cSHong Zhang   YdotRHS_fast = ((TS_RK*)rk->subts_fast->data)->YdotRHS;
333*bb42530cSHong Zhang   YdotRHS_slow = ((TS_RK*)rk->subts_slow->data)->YdotRHS;
334474dd773SHong Zhang   for (i=0; i<s; i++) {
335474dd773SHong Zhang     ierr = VecGetSubVector(YdotRHS[i],rk->is_slow,&YdotRHS_slow[i]);CHKERRQ(ierr);
336474dd773SHong Zhang     ierr = VecGetSubVector(YdotRHS[i],rk->is_fast,&YdotRHS_fast[i]);CHKERRQ(ierr);
337474dd773SHong Zhang   }
338474dd773SHong Zhang   ierr = VecCopy(ts->vec_sol,rk->X0);CHKERRQ(ierr);
339474dd773SHong Zhang   /* propogate both slow and fast components using large time steps */
340474dd773SHong Zhang   for (i=0; i<s; i++) {
341474dd773SHong Zhang     rk->stage_time = t + h*c[i];
342474dd773SHong Zhang     ierr = TSPreStage(ts,rk->stage_time);CHKERRQ(ierr);
343474dd773SHong Zhang     ierr = VecCopy(ts->vec_sol,Y[i]);CHKERRQ(ierr);
344474dd773SHong Zhang     /*ierr = VecCopy(ts->vec_sol,Yc[i]);CHKERRQ(ierr);*/
345474dd773SHong Zhang     ierr = VecGetSubVector(Y[i],rk->is_fast,&Yfast);CHKERRQ(ierr);
346*bb42530cSHong Zhang     ierr = VecGetSubVector(Y[i],rk->is_slow,&Yslow);CHKERRQ(ierr);
347474dd773SHong Zhang     for (j=0; j<i; j++) w[j] = h*A[i*s+j];
348474dd773SHong Zhang     ierr = VecMAXPY(Yfast,i,w,YdotRHS_fast);CHKERRQ(ierr);
349*bb42530cSHong Zhang     ierr = VecMAXPY(Yslow,i,w,YdotRHS_slow);CHKERRQ(ierr);
350474dd773SHong Zhang     ierr = VecRestoreSubVector(Y[i],rk->is_fast,&Yfast);CHKERRQ(ierr);
351*bb42530cSHong Zhang     ierr = VecRestoreSubVector(Y[i],rk->is_slow,&Yslow);CHKERRQ(ierr);
352474dd773SHong Zhang     ierr = TSPostStage(ts,rk->stage_time,i,Y); CHKERRQ(ierr);
353474dd773SHong Zhang     ierr = TSComputeRHSFunction(rk->subts_slow,t+h*c[i],Y[i],YdotRHS_slow[i]);CHKERRQ(ierr);
354474dd773SHong Zhang     ierr = TSComputeRHSFunction(rk->subts_fast,t+h*c[i],Y[i],YdotRHS_fast[i]);CHKERRQ(ierr);
355474dd773SHong Zhang   }
356474dd773SHong Zhang   rk->slow = PETSC_TRUE;
357*bb42530cSHong Zhang   /* update the value of slow components when using slow time step */
358474dd773SHong Zhang   ierr = TSEvaluateStep_MRKSPLIT(ts,tab->order,ts->vec_sol,NULL);CHKERRQ(ierr);
359*bb42530cSHong Zhang 
360*bb42530cSHong Zhang   if (((TS_RK*)rk->subts_fast->data)->subts_fast) {
361*bb42530cSHong Zhang     rk->subts_fast->ptime = t;
362*bb42530cSHong Zhang     rk->subts_fast->time_step = h/rk->dtratio;
363*bb42530cSHong Zhang     for (k=0; k<rk->dtratio; k++) {
364*bb42530cSHong Zhang       ierr = TSStep_MRKSPLIT(rk->subts_fast);CHKERRQ(ierr);
365*bb42530cSHong Zhang     }
366*bb42530cSHong Zhang   } else {
367474dd773SHong Zhang     for (k=0; k<rk->dtratio; k++) {
368474dd773SHong Zhang       /* propogate fast component using small time steps */
369474dd773SHong Zhang       for (i=0; i<s; i++) {
370474dd773SHong Zhang         rk->stage_time = t + h/rk->dtratio*c[i];
371474dd773SHong Zhang         ierr = TSPreStage(ts,rk->stage_time);CHKERRQ(ierr);
372474dd773SHong Zhang         ierr = VecCopy(ts->vec_sol,Y[i]);CHKERRQ(ierr);
373474dd773SHong Zhang         /* stage value for fast components */
374474dd773SHong Zhang         for (j=0; j<i; j++) w[j] = h/rk->dtratio*A[i*s+j];
375474dd773SHong Zhang         ierr = VecGetSubVector(Y[i],rk->is_fast,&Yfast);CHKERRQ(ierr);
376474dd773SHong Zhang         ierr = VecMAXPY(Yfast,i,w,YdotRHS_fast);CHKERRQ(ierr);
377474dd773SHong Zhang         ierr = VecRestoreSubVector(Y[i],rk->is_fast,&Yfast);CHKERRQ(ierr);
378474dd773SHong Zhang         /* stage value for slow components */
379474dd773SHong Zhang         ierr = VecGetSubVector(Y[i],rk->is_slow,&Yslow);CHKERRQ(ierr);
380474dd773SHong Zhang         ierr = TSInterpolate_MRKSPLIT(ts,t+k*h/rk->dtratio+h/rk->dtratio*c[i],Yslow);CHKERRQ(ierr);
381474dd773SHong Zhang         ierr = VecRestoreSubVector(Y[i],rk->is_slow,&Yslow);CHKERRQ(ierr);
382474dd773SHong Zhang         ierr = TSPostStage(ts,rk->stage_time,i,Y);CHKERRQ(ierr);
383474dd773SHong Zhang         /* compute the stage RHS for fast components */
384*bb42530cSHong Zhang         ierr = TSComputeRHSFunction(rk->subts_fast,t+k*h*rk->dtratio+h/rk->dtratio*c[i],Y[i],YdotRHS_fast[i]);CHKERRQ(ierr);
385474dd773SHong Zhang       }
386474dd773SHong Zhang       /* update the value of fast components when using fast time step */
387474dd773SHong Zhang       rk->slow = PETSC_FALSE;
388474dd773SHong Zhang       ierr = TSEvaluateStep_MRKSPLIT(ts,tab->order,ts->vec_sol,NULL);CHKERRQ(ierr);
389474dd773SHong Zhang     }
390*bb42530cSHong Zhang   }
391474dd773SHong Zhang   for (i=0; i<s; i++) {
392474dd773SHong Zhang     ierr = VecRestoreSubVector(YdotRHS[i],rk->is_slow,&YdotRHS_slow[i]);CHKERRQ(ierr);
393474dd773SHong Zhang     ierr = VecRestoreSubVector(YdotRHS[i],rk->is_fast,&YdotRHS_fast[i]);CHKERRQ(ierr);
394474dd773SHong Zhang   }
395*bb42530cSHong Zhang   ts->ptime = t + ts->time_step;
396474dd773SHong Zhang   ts->time_step = next_time_step;
397474dd773SHong Zhang   rk->status = TS_STEP_COMPLETE;
398474dd773SHong Zhang   PetscFunctionReturn(0);
399474dd773SHong Zhang }
400474dd773SHong Zhang 
401474dd773SHong Zhang /*@C
402474dd773SHong Zhang   TSRKSetMultirateType - Set the type of RK Multirate scheme
403474dd773SHong Zhang 
404474dd773SHong Zhang   Logically collective
405474dd773SHong Zhang 
406474dd773SHong Zhang   Input Parameter:
407474dd773SHong Zhang +  ts - timestepping context
408474dd773SHong Zhang -  mrktype - type of MRK-scheme
409474dd773SHong Zhang 
410474dd773SHong Zhang   Options Database:
411474dd773SHong Zhang .   -ts_rk_multiarte_type - <none,nonsplit,split>
412474dd773SHong Zhang 
413474dd773SHong Zhang   Level: intermediate
414474dd773SHong Zhang @*/
415474dd773SHong Zhang PetscErrorCode TSRKSetMultirateType(TS ts, TSMRKType mrktype)
416474dd773SHong Zhang {
417474dd773SHong Zhang   TS_RK          *rk = (TS_RK*)ts->data;
418474dd773SHong Zhang   PetscErrorCode ierr;
419474dd773SHong Zhang 
420474dd773SHong Zhang   PetscFunctionBegin;
421474dd773SHong Zhang   PetscValidHeaderSpecific(ts,TS_CLASSID,1);
422474dd773SHong Zhang   switch(mrktype){
423474dd773SHong Zhang     case TSMRKNONE:
424474dd773SHong Zhang       break;
425474dd773SHong Zhang     case TSMRKNONSPLIT:
426474dd773SHong Zhang       ts->ops->step           = TSStep_MRKNONSPLIT;
427474dd773SHong Zhang       ts->ops->interpolate    = TSInterpolate_MRKNONSPLIT;
428474dd773SHong Zhang       rk->dtratio             = 2;
429474dd773SHong Zhang       ierr = TSRKSetType(ts,TSMRKDefault);CHKERRQ(ierr);
430474dd773SHong Zhang       ierr = PetscObjectComposeFunction((PetscObject)ts,"TSSetUp_MRKNONSPLIT_C",TSSetUp_MRKNONSPLIT);CHKERRQ(ierr);
431474dd773SHong Zhang       ierr = PetscObjectComposeFunction((PetscObject)ts,"TSReset_MRKNONSPLIT_C",TSReset_MRKNONSPLIT);CHKERRQ(ierr);
432474dd773SHong Zhang       break;
433474dd773SHong Zhang     case TSMRKSPLIT:
434474dd773SHong Zhang       ts->ops->step           = TSStep_MRKSPLIT;
435474dd773SHong Zhang       ts->ops->evaluatestep   = TSEvaluateStep_MRKSPLIT;
436474dd773SHong Zhang       ts->ops->interpolate    = TSInterpolate_MRKSPLIT;
437474dd773SHong Zhang       rk->dtratio             = 2;
438474dd773SHong Zhang       ierr = TSRKSetType(ts,TSMRKDefault);CHKERRQ(ierr);
439474dd773SHong Zhang       ierr = PetscObjectComposeFunction((PetscObject)ts,"TSSetUp_MRKSPLIT_C",TSSetUp_MRKSPLIT);CHKERRQ(ierr);
440474dd773SHong Zhang       ierr = PetscObjectComposeFunction((PetscObject)ts,"TSReset_MRKSPLIT_C",TSReset_MRKSPLIT);CHKERRQ(ierr);
441474dd773SHong Zhang       break;
442474dd773SHong Zhang     default :
443474dd773SHong Zhang       SETERRQ1(PetscObjectComm((PetscObject)ts),PETSC_ERR_ARG_UNKNOWN_TYPE,"Unknown type '%s'",mrktype);
444474dd773SHong Zhang   }
445474dd773SHong Zhang   PetscFunctionReturn(0);
446474dd773SHong Zhang }
447474dd773SHong Zhang 
448