xref: /petsc/src/ts/impls/explicit/rk/rk.c (revision d190556429ad48cdb35b176ba6712769ba23bffc)
1 /*
2   Code for time stepping with the Runge-Kutta method
3 
4   Notes:
5   The general system is written as
6 
7   Udot = F(t,U)
8 
9 */
10 #include <petsc/private/tsimpl.h>                /*I   "petscts.h"   I*/
11 #include <petscdm.h>
12 
13 static TSRKType  TSRKDefault = TSRK3BS;
14 static PetscBool TSRKRegisterAllCalled;
15 static PetscBool TSRKPackageInitialized;
16 
17 typedef struct _RKTableau *RKTableau;
18 struct _RKTableau {
19   char      *name;
20   PetscInt   order;               /* Classical approximation order of the method i              */
21   PetscInt   s;                   /* Number of stages                                           */
22   PetscInt   p;                   /* Interpolation order                                        */
23   PetscBool  FSAL;                /* flag to indicate if tableau is FSAL                        */
24   PetscReal *A,*b,*c;             /* Tableau                                                    */
25   PetscReal *bembed;              /* Embedded formula of order one less (order-1)               */
26   PetscReal *binterp;             /* Dense output formula                                       */
27   PetscReal  ccfl;                /* Placeholder for CFL coefficient relative to forward Euler  */
28 };
29 typedef struct _RKTableauLink *RKTableauLink;
30 struct _RKTableauLink {
31   struct _RKTableau tab;
32   RKTableauLink     next;
33 };
34 static RKTableauLink RKTableauList;
35 
36 typedef struct {
37   RKTableau    tableau;
38   Vec          *Y;               /* States computed during the step */
39   Vec          *YdotRHS;         /* Function evaluations for the non-stiff part */
40   Vec          *VecDeltaLam;     /* Increment of the adjoint sensitivity w.r.t IC at stage */
41   Vec          *VecDeltaMu;      /* Increment of the adjoint sensitivity w.r.t P at stage */
42   Vec          *VecSensiTemp;    /* Vector to be timed with Jacobian transpose */
43   Vec          VecCostIntegral0; /* backup for roll-backs due to events */
44   PetscScalar  *work;            /* Scalar work */
45   PetscReal    stage_time;
46   TSStepStatus status;
47   PetscReal    ptime;
48   PetscReal    time_step;
49 } TS_RK;
50 
51 /*MC
52      TSRK1 - First order forward Euler scheme.
53 
54      This method has one stage.
55 
56      Level: advanced
57 
58 .seealso: TSRK
59 M*/
60 /*MC
61      TSRK2A - Second order RK scheme.
62 
63      This method has two stages.
64 
65      Level: advanced
66 
67 .seealso: TSRK
68 M*/
69 /*MC
70      TSRK3 - Third order RK scheme.
71 
72      This method has three stages.
73 
74      Level: advanced
75 
76 .seealso: TSRK
77 M*/
78 /*MC
79      TSRK3BS - Third order RK scheme of Bogacki-Shampine with 2nd order embedded method.
80 
81      This method has four stages with the First Same As Last (FSAL) property.
82 
83      Level: advanced
84 
85      Notes: https://doi.org/10.1016/0893-9659(89)90079-7
86 
87 .seealso: TSRK
88 M*/
89 /*MC
90      TSRK4 - Fourth order RK scheme.
91 
92      This is the classical Runge-Kutta method with four stages.
93 
94      Level: advanced
95 
96 .seealso: TSRK
97 M*/
98 /*MC
99      TSRK5F - Fifth order Fehlberg RK scheme with a 4th order embedded method.
100 
101      This method has six stages.
102 
103      Level: advanced
104 
105 .seealso: TSRK
106 M*/
107 /*MC
108      TSRK5DP - Fifth order Dormand-Prince RK scheme with the 4th order embedded method.
109 
110      This method has seven stages with the First Same As Last (FSAL) property.
111 
112      Level: advanced
113 
114      Notes: https://doi.org/10.1016/0771-050X(80)90013-3
115 
116 .seealso: TSRK
117 M*/
118 
119 /*@C
120   TSRKRegisterAll - Registers all of the Runge-Kutta explicit methods in TSRK
121 
122   Not Collective, but should be called by all processes which will need the schemes to be registered
123 
124   Level: advanced
125 
126 .keywords: TS, TSRK, register, all
127 
128 .seealso:  TSRKRegisterDestroy()
129 @*/
130 PetscErrorCode TSRKRegisterAll(void)
131 {
132   PetscErrorCode ierr;
133 
134   PetscFunctionBegin;
135   if (TSRKRegisterAllCalled) PetscFunctionReturn(0);
136   TSRKRegisterAllCalled = PETSC_TRUE;
137 
138   {
139     const PetscReal
140       A[1][1] = {{0.0}},
141       b[1]    = {1.0};
142     ierr = TSRKRegister(TSRK1FE,1,1,&A[0][0],b,NULL,NULL,0,NULL);CHKERRQ(ierr);
143   }
144   {
145     const PetscReal
146       A[2][2]     = {{0.0,0.0},
147                     {1.0,0.0}},
148       b[2]        = {0.5,0.5},
149       bembed[2]   = {1.0,0};
150     ierr = TSRKRegister(TSRK2A,2,2,&A[0][0],b,NULL,bembed,0,NULL);CHKERRQ(ierr);
151   }
152   {
153     const PetscReal
154       A[3][3] = {{0,0,0},
155                  {2.0/3.0,0,0},
156                  {-1.0/3.0,1.0,0}},
157       b[3]    = {0.25,0.5,0.25};
158     ierr = TSRKRegister(TSRK3,3,3,&A[0][0],b,NULL,NULL,0,NULL);CHKERRQ(ierr);
159   }
160   {
161     const PetscReal
162       A[4][4] = {{0,0,0,0},
163                  {1.0/2.0,0,0,0},
164                  {0,3.0/4.0,0,0},
165                  {2.0/9.0,1.0/3.0,4.0/9.0,0}},
166       b[4]    = {2.0/9.0,1.0/3.0,4.0/9.0,0},
167       bembed[4] = {7.0/24.0,1.0/4.0,1.0/3.0,1.0/8.0};
168     ierr = TSRKRegister(TSRK3BS,3,4,&A[0][0],b,NULL,bembed,0,NULL);CHKERRQ(ierr);
169   }
170   {
171     const PetscReal
172       A[4][4] = {{0,0,0,0},
173                  {0.5,0,0,0},
174                  {0,0.5,0,0},
175                  {0,0,1.0,0}},
176       b[4]    = {1.0/6.0,1.0/3.0,1.0/3.0,1.0/6.0};
177     ierr = TSRKRegister(TSRK4,4,4,&A[0][0],b,NULL,NULL,0,NULL);CHKERRQ(ierr);
178   }
179   {
180     const PetscReal
181       A[6][6]   = {{0,0,0,0,0,0},
182                    {0.25,0,0,0,0,0},
183                    {3.0/32.0,9.0/32.0,0,0,0,0},
184                    {1932.0/2197.0,-7200.0/2197.0,7296.0/2197.0,0,0,0},
185                    {439.0/216.0,-8.0,3680.0/513.0,-845.0/4104.0,0,0},
186                    {-8.0/27.0,2.0,-3544.0/2565.0,1859.0/4104.0,-11.0/40.0,0}},
187       b[6]      = {16.0/135.0,0,6656.0/12825.0,28561.0/56430.0,-9.0/50.0,2.0/55.0},
188       bembed[6] = {25.0/216.0,0,1408.0/2565.0,2197.0/4104.0,-1.0/5.0,0};
189     ierr = TSRKRegister(TSRK5F,5,6,&A[0][0],b,NULL,bembed,0,NULL);CHKERRQ(ierr);
190   }
191   {
192     const PetscReal
193       A[7][7]   = {{0,0,0,0,0,0,0},
194                    {1.0/5.0,0,0,0,0,0,0},
195                    {3.0/40.0,9.0/40.0,0,0,0,0,0},
196                    {44.0/45.0,-56.0/15.0,32.0/9.0,0,0,0,0},
197                    {19372.0/6561.0,-25360.0/2187.0,64448.0/6561.0,-212.0/729.0,0,0,0},
198                    {9017.0/3168.0,-355.0/33.0,46732.0/5247.0,49.0/176.0,-5103.0/18656.0,0,0},
199                    {35.0/384.0,0,500.0/1113.0,125.0/192.0,-2187.0/6784.0,11.0/84.0,0}},
200       b[7]      = {35.0/384.0,0,500.0/1113.0,125.0/192.0,-2187.0/6784.0,11.0/84.0,0},
201       bembed[7] = {5179.0/57600.0,0,7571.0/16695.0,393.0/640.0,-92097.0/339200.0,187.0/2100.0,1.0/40.0};
202     ierr = TSRKRegister(TSRK5DP,5,7,&A[0][0],b,NULL,bembed,0,NULL);CHKERRQ(ierr);
203   }
204   PetscFunctionReturn(0);
205 }
206 
207 /*@C
208    TSRKRegisterDestroy - Frees the list of schemes that were registered by TSRKRegister().
209 
210    Not Collective
211 
212    Level: advanced
213 
214 .keywords: TSRK, register, destroy
215 .seealso: TSRKRegister(), TSRKRegisterAll()
216 @*/
217 PetscErrorCode TSRKRegisterDestroy(void)
218 {
219   PetscErrorCode ierr;
220   RKTableauLink link;
221 
222   PetscFunctionBegin;
223   while ((link = RKTableauList)) {
224     RKTableau t = &link->tab;
225     RKTableauList = link->next;
226     ierr = PetscFree3(t->A,t->b,t->c);  CHKERRQ(ierr);
227     ierr = PetscFree (t->bembed);       CHKERRQ(ierr);
228     ierr = PetscFree (t->binterp);      CHKERRQ(ierr);
229     ierr = PetscFree (t->name);         CHKERRQ(ierr);
230     ierr = PetscFree (link);            CHKERRQ(ierr);
231   }
232   TSRKRegisterAllCalled = PETSC_FALSE;
233   PetscFunctionReturn(0);
234 }
235 
236 /*@C
237   TSRKInitializePackage - This function initializes everything in the TSRK package. It is called
238   from PetscDLLibraryRegister() when using dynamic libraries, and on the first call to TSCreate_RK()
239   when using static libraries.
240 
241   Level: developer
242 
243 .keywords: TS, TSRK, initialize, package
244 .seealso: PetscInitialize()
245 @*/
246 PetscErrorCode TSRKInitializePackage(void)
247 {
248   PetscErrorCode ierr;
249 
250   PetscFunctionBegin;
251   if (TSRKPackageInitialized) PetscFunctionReturn(0);
252   TSRKPackageInitialized = PETSC_TRUE;
253   ierr = TSRKRegisterAll();CHKERRQ(ierr);
254   ierr = PetscRegisterFinalize(TSRKFinalizePackage);CHKERRQ(ierr);
255   PetscFunctionReturn(0);
256 }
257 
258 /*@C
259   TSRKFinalizePackage - This function destroys everything in the TSRK package. It is
260   called from PetscFinalize().
261 
262   Level: developer
263 
264 .keywords: Petsc, destroy, package
265 .seealso: PetscFinalize()
266 @*/
267 PetscErrorCode TSRKFinalizePackage(void)
268 {
269   PetscErrorCode ierr;
270 
271   PetscFunctionBegin;
272   TSRKPackageInitialized = PETSC_FALSE;
273   ierr = TSRKRegisterDestroy();CHKERRQ(ierr);
274   PetscFunctionReturn(0);
275 }
276 
277 /*@C
278    TSRKRegister - register an RK scheme by providing the entries in the Butcher tableau and optionally embedded approximations and interpolation
279 
280    Not Collective, but the same schemes should be registered on all processes on which they will be used
281 
282    Input Parameters:
283 +  name - identifier for method
284 .  order - approximation order of method
285 .  s - number of stages, this is the dimension of the matrices below
286 .  A - stage coefficients (dimension s*s, row-major)
287 .  b - step completion table (dimension s; NULL to use last row of A)
288 .  c - abscissa (dimension s; NULL to use row sums of A)
289 .  bembed - completion table for embedded method (dimension s; NULL if not available)
290 .  p - Order of the interpolation scheme, equal to the number of columns of binterp
291 -  binterp - Coefficients of the interpolation formula (dimension s*p; NULL to reuse b with p=1)
292 
293    Notes:
294    Several RK methods are provided, this function is only needed to create new methods.
295 
296    Level: advanced
297 
298 .keywords: TS, register
299 
300 .seealso: TSRK
301 @*/
302 PetscErrorCode TSRKRegister(TSRKType name,PetscInt order,PetscInt s,
303                             const PetscReal A[],const PetscReal b[],const PetscReal c[],
304                             const PetscReal bembed[],PetscInt p,const PetscReal binterp[])
305 {
306   PetscErrorCode  ierr;
307   RKTableauLink   link;
308   RKTableau       t;
309   PetscInt        i,j;
310 
311   PetscFunctionBegin;
312   PetscValidCharPointer(name,1);
313   PetscValidRealPointer(A,4);
314   if (b) PetscValidRealPointer(b,5);
315   if (c) PetscValidRealPointer(c,6);
316   if (bembed) PetscValidRealPointer(bembed,7);
317   if (binterp || p > 1) PetscValidRealPointer(binterp,9);
318 
319   ierr = PetscNew(&link);CHKERRQ(ierr);
320   t = &link->tab;
321 
322   ierr = PetscStrallocpy(name,&t->name);CHKERRQ(ierr);
323   t->order = order;
324   t->s = s;
325   ierr = PetscMalloc3(s*s,&t->A,s,&t->b,s,&t->c);CHKERRQ(ierr);
326   ierr = PetscMemcpy(t->A,A,s*s*sizeof(A[0]));CHKERRQ(ierr);
327   if (b)  { ierr = PetscMemcpy(t->b,b,s*sizeof(b[0]));CHKERRQ(ierr); }
328   else for (i=0; i<s; i++) t->b[i] = A[(s-1)*s+i];
329   if (c)  { ierr = PetscMemcpy(t->c,c,s*sizeof(c[0]));CHKERRQ(ierr); }
330   else for (i=0; i<s; i++) for (j=0,t->c[i]=0; j<s; j++) t->c[i] += A[i*s+j];
331   t->FSAL = PETSC_TRUE;
332   for (i=0; i<s; i++) if (t->A[(s-1)*s+i] != t->b[i]) t->FSAL = PETSC_FALSE;
333 
334   if (bembed) {
335     ierr = PetscMalloc1(s,&t->bembed);CHKERRQ(ierr);
336     ierr = PetscMemcpy(t->bembed,bembed,s*sizeof(bembed[0]));CHKERRQ(ierr);
337   }
338 
339   if (!binterp) { p = 1; binterp = t->b; }
340   t->p = p;
341   ierr = PetscMalloc1(s*p,&t->binterp);CHKERRQ(ierr);
342   ierr = PetscMemcpy(t->binterp,binterp,s*p*sizeof(binterp[0]));CHKERRQ(ierr);
343 
344   link->next = RKTableauList;
345   RKTableauList = link;
346   PetscFunctionReturn(0);
347 }
348 
349 /*
350  The step completion formula is
351 
352  x1 = x0 + h b^T YdotRHS
353 
354  This function can be called before or after ts->vec_sol has been updated.
355  Suppose we have a completion formula (b) and an embedded formula (be) of different order.
356  We can write
357 
358  x1e = x0 + h be^T YdotRHS
359      = x1 - h b^T YdotRHS + h be^T YdotRHS
360      = x1 + h (be - b)^T YdotRHS
361 
362  so we can evaluate the method with different order even after the step has been optimistically completed.
363 */
364 static PetscErrorCode TSEvaluateStep_RK(TS ts,PetscInt order,Vec X,PetscBool *done)
365 {
366   TS_RK         *rk   = (TS_RK*)ts->data;
367   RKTableau      tab  = rk->tableau;
368   PetscScalar   *w    = rk->work;
369   PetscReal      h;
370   PetscInt       s    = tab->s,j;
371   PetscErrorCode ierr;
372 
373   PetscFunctionBegin;
374   switch (rk->status) {
375   case TS_STEP_INCOMPLETE:
376   case TS_STEP_PENDING:
377     h = ts->time_step; break;
378   case TS_STEP_COMPLETE:
379     h = ts->ptime - ts->ptime_prev; break;
380   default: SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_PLIB,"Invalid TSStepStatus");
381   }
382   if (order == tab->order) {
383     if (rk->status == TS_STEP_INCOMPLETE) {
384       ierr = VecCopy(ts->vec_sol,X);CHKERRQ(ierr);
385       for (j=0; j<s; j++) w[j] = h*tab->b[j];
386       ierr = VecMAXPY(X,s,w,rk->YdotRHS);CHKERRQ(ierr);
387     } else {ierr = VecCopy(ts->vec_sol,X);CHKERRQ(ierr);}
388     PetscFunctionReturn(0);
389   } else if (order == tab->order-1) {
390     if (!tab->bembed) goto unavailable;
391     if (rk->status == TS_STEP_INCOMPLETE) { /* Complete with the embedded method (be) */
392       ierr = VecCopy(ts->vec_sol,X);CHKERRQ(ierr);
393       for (j=0; j<s; j++) w[j] = h*tab->bembed[j];
394       ierr = VecMAXPY(X,s,w,rk->YdotRHS);CHKERRQ(ierr);
395     } else { /* Rollback and re-complete using (be-b) */
396       ierr = VecCopy(ts->vec_sol,X);CHKERRQ(ierr);
397       for (j=0; j<s; j++) w[j] = h*(tab->bembed[j] - tab->b[j]);
398       ierr = VecMAXPY(X,s,w,rk->YdotRHS);CHKERRQ(ierr);
399       if (ts->vec_costintegral && ts->costintegralfwd) {
400         ierr = VecCopy(rk->VecCostIntegral0,ts->vec_costintegral);CHKERRQ(ierr);
401       }
402     }
403     if (done) *done = PETSC_TRUE;
404     PetscFunctionReturn(0);
405   }
406 unavailable:
407   if (done) *done = PETSC_FALSE;
408   else SETERRQ3(PetscObjectComm((PetscObject)ts),PETSC_ERR_SUP,"RK '%s' of order %D cannot evaluate step at order %D. Consider using -ts_adapt_type none or a different method that has an embedded estimate.",tab->name,tab->order,order);
409   PetscFunctionReturn(0);
410 }
411 
412 static PetscErrorCode TSForwardCostIntegral_RK(TS ts)
413 {
414   TS_RK           *rk = (TS_RK*)ts->data;
415   RKTableau       tab = rk->tableau;
416   const PetscInt  s = tab->s;
417   const PetscReal *b = tab->b,*c = tab->c;
418   Vec             *Y = rk->Y;
419   PetscInt        i;
420   PetscErrorCode  ierr;
421 
422   PetscFunctionBegin;
423   /* backup cost integral */
424   ierr = VecCopy(ts->vec_costintegral,rk->VecCostIntegral0);CHKERRQ(ierr);
425   for (i=s-1; i>=0; i--) {
426     /* Evolve ts->vec_costintegral to compute integrals */
427     ierr = TSAdjointComputeCostIntegrand(ts,rk->ptime+rk->time_step*(1.0-c[i]),Y[i],ts->vec_costintegrand);CHKERRQ(ierr);
428     ierr = VecAXPY(ts->vec_costintegral,rk->time_step*b[i],ts->vec_costintegrand);CHKERRQ(ierr);
429   }
430   PetscFunctionReturn(0);
431 }
432 
433 static PetscErrorCode TSAdjointCostIntegral_RK(TS ts)
434 {
435   TS_RK           *rk = (TS_RK*)ts->data;
436   RKTableau       tab = rk->tableau;
437   const PetscInt  s = tab->s;
438   const PetscReal *b = tab->b,*c = tab->c;
439   Vec             *Y = rk->Y;
440   PetscInt        i;
441   PetscErrorCode  ierr;
442 
443   PetscFunctionBegin;
444   for (i=s-1; i>=0; i--) {
445     /* Evolve ts->vec_costintegral to compute integrals */
446     ierr = TSAdjointComputeCostIntegrand(ts,ts->ptime-ts->time_step*(1.0-c[i]),Y[i],ts->vec_costintegrand);CHKERRQ(ierr);
447     ierr = VecAXPY(ts->vec_costintegral,-ts->time_step*b[i],ts->vec_costintegrand);CHKERRQ(ierr);
448   }
449   PetscFunctionReturn(0);
450 }
451 
452 static PetscErrorCode TSRollBack_RK(TS ts)
453 {
454   TS_RK           *rk = (TS_RK*)ts->data;
455   RKTableau       tab = rk->tableau;
456   const PetscInt  s  = tab->s;
457   const PetscReal *b = tab->b;
458   PetscScalar     *w = rk->work;
459   Vec             *YdotRHS = rk->YdotRHS;
460   PetscInt        j;
461   PetscReal       h;
462   PetscErrorCode  ierr;
463 
464   PetscFunctionBegin;
465   switch (rk->status) {
466   case TS_STEP_INCOMPLETE:
467   case TS_STEP_PENDING:
468     h = ts->time_step; break;
469   case TS_STEP_COMPLETE:
470     h = ts->ptime - ts->ptime_prev; break;
471   default: SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_PLIB,"Invalid TSStepStatus");
472   }
473   for (j=0; j<s; j++) w[j] = -h*b[j];
474   ierr = VecMAXPY(ts->vec_sol,s,w,YdotRHS);CHKERRQ(ierr);
475   PetscFunctionReturn(0);
476 }
477 
478 static PetscErrorCode TSStep_RK(TS ts)
479 {
480   TS_RK           *rk   = (TS_RK*)ts->data;
481   RKTableau        tab  = rk->tableau;
482   const PetscInt   s = tab->s;
483   const PetscReal *A = tab->A,*c = tab->c;
484   PetscScalar     *w = rk->work;
485   Vec             *Y = rk->Y,*YdotRHS = rk->YdotRHS;
486   PetscBool        FSAL = tab->FSAL;
487   TSAdapt          adapt;
488   PetscInt         i,j;
489   PetscInt         rejections = 0;
490   PetscBool        stageok,accept = PETSC_TRUE;
491   PetscReal        next_time_step = ts->time_step;
492   PetscErrorCode   ierr;
493 
494   PetscFunctionBegin;
495   if (ts->steprollback || ts->steprestart) FSAL = PETSC_FALSE;
496   if (FSAL) { ierr = VecCopy(YdotRHS[s-1],YdotRHS[0]);CHKERRQ(ierr); }
497 
498   rk->status = TS_STEP_INCOMPLETE;
499   while (!ts->reason && rk->status != TS_STEP_COMPLETE) {
500     PetscReal t = ts->ptime;
501     PetscReal h = ts->time_step;
502     for (i=0; i<s; i++) {
503       if (FSAL && !i) continue;
504       rk->stage_time = t + h*c[i];
505       ierr = TSPreStage(ts,rk->stage_time); CHKERRQ(ierr);
506       ierr = VecCopy(ts->vec_sol,Y[i]);CHKERRQ(ierr);
507       for (j=0; j<i; j++) w[j] = h*A[i*s+j];
508       ierr = VecMAXPY(Y[i],i,w,YdotRHS);CHKERRQ(ierr);
509       ierr = TSPostStage(ts,rk->stage_time,i,Y); CHKERRQ(ierr);
510       ierr = TSGetAdapt(ts,&adapt);CHKERRQ(ierr);
511       ierr = TSAdaptCheckStage(adapt,ts,rk->stage_time,Y[i],&stageok);CHKERRQ(ierr);
512       if (!stageok) goto reject_step;
513       ierr = TSComputeRHSFunction(ts,t+h*c[i],Y[i],YdotRHS[i]);CHKERRQ(ierr);
514     }
515 
516     rk->status = TS_STEP_INCOMPLETE;
517     ierr = TSEvaluateStep(ts,tab->order,ts->vec_sol,NULL);CHKERRQ(ierr);
518     rk->status = TS_STEP_PENDING;
519     ierr = TSGetAdapt(ts,&adapt);CHKERRQ(ierr);
520     ierr = TSAdaptCandidatesClear(adapt);CHKERRQ(ierr);
521     ierr = TSAdaptCandidateAdd(adapt,tab->name,tab->order,1,tab->ccfl,1.*tab->s,PETSC_TRUE);CHKERRQ(ierr);
522     ierr = TSAdaptChoose(adapt,ts,ts->time_step,NULL,&next_time_step,&accept);CHKERRQ(ierr);
523     rk->status = accept ? TS_STEP_COMPLETE : TS_STEP_INCOMPLETE;
524     if (!accept) { /* Roll back the current step */
525       ierr = TSRollBack_RK(ts);CHKERRQ(ierr);
526       ts->time_step = next_time_step;
527       goto reject_step;
528     }
529 
530     if (ts->costintegralfwd) { /* Save the info for the later use in cost integral evaluation*/
531       rk->ptime     = ts->ptime;
532       rk->time_step = ts->time_step;
533     }
534 
535     ts->ptime += ts->time_step;
536     ts->time_step = next_time_step;
537     break;
538 
539   reject_step:
540     ts->reject++; accept = PETSC_FALSE;
541     if (!ts->reason && ++rejections > ts->max_reject && ts->max_reject >= 0) {
542       ts->reason = TS_DIVERGED_STEP_REJECTED;
543       ierr = PetscInfo2(ts,"Step=%D, step rejections %D greater than current TS allowed, stopping solve\n",ts->steps,rejections);CHKERRQ(ierr);
544     }
545   }
546   PetscFunctionReturn(0);
547 }
548 
549 static PetscErrorCode TSAdjointSetUp_RK(TS ts)
550 {
551   TS_RK         *rk  = (TS_RK*)ts->data;
552   RKTableau      tab = rk->tableau;
553   PetscInt       s   = tab->s;
554   PetscErrorCode ierr;
555 
556   PetscFunctionBegin;
557   if (ts->adjointsetupcalled++) PetscFunctionReturn(0);
558   ierr = VecDuplicateVecs(ts->vecs_sensi[0],s*ts->numcost,&rk->VecDeltaLam);CHKERRQ(ierr);
559   if(ts->vecs_sensip) {
560     ierr = VecDuplicateVecs(ts->vecs_sensip[0],s*ts->numcost,&rk->VecDeltaMu);CHKERRQ(ierr);
561   }
562   ierr = VecDuplicateVecs(ts->vecs_sensi[0],ts->numcost,&rk->VecSensiTemp);CHKERRQ(ierr);
563   PetscFunctionReturn(0);
564 }
565 
566 static PetscErrorCode TSAdjointStep_RK(TS ts)
567 {
568   TS_RK           *rk   = (TS_RK*)ts->data;
569   RKTableau        tab  = rk->tableau;
570   const PetscInt   s    = tab->s;
571   const PetscReal *A = tab->A,*b = tab->b,*c = tab->c;
572   PetscScalar     *w    = rk->work;
573   Vec             *Y    = rk->Y,*VecDeltaLam = rk->VecDeltaLam,*VecDeltaMu = rk->VecDeltaMu,*VecSensiTemp = rk->VecSensiTemp;
574   PetscInt         i,j,nadj;
575   PetscReal        t = ts->ptime;
576   PetscErrorCode   ierr;
577   PetscReal        h = ts->time_step;
578   Mat              J,Jp;
579 
580   PetscFunctionBegin;
581   rk->status = TS_STEP_INCOMPLETE;
582   for (i=s-1; i>=0; i--) {
583     rk->stage_time = t + h*(1.0-c[i]);
584     for (nadj=0; nadj<ts->numcost; nadj++) {
585       ierr = VecCopy(ts->vecs_sensi[nadj],VecSensiTemp[nadj]);CHKERRQ(ierr);
586       ierr = VecScale(VecSensiTemp[nadj],-h*b[i]);CHKERRQ(ierr);
587       for (j=i+1; j<s; j++) {
588         ierr = VecAXPY(VecSensiTemp[nadj],-h*A[j*s+i],VecDeltaLam[nadj*s+j]);CHKERRQ(ierr);
589       }
590     }
591     /* Stage values of lambda */
592     ierr = TSGetRHSJacobian(ts,&J,&Jp,NULL,NULL);CHKERRQ(ierr);
593     ierr = TSComputeRHSJacobian(ts,rk->stage_time,Y[i],J,Jp);CHKERRQ(ierr);
594     if (ts->vec_costintegral) {
595       ierr = TSAdjointComputeDRDYFunction(ts,rk->stage_time,Y[i],ts->vecs_drdy);CHKERRQ(ierr);
596     }
597     for (nadj=0; nadj<ts->numcost; nadj++) {
598       ierr = MatMultTranspose(J,VecSensiTemp[nadj],VecDeltaLam[nadj*s+i]);CHKERRQ(ierr);
599       if (ts->vec_costintegral) {
600         ierr = VecAXPY(VecDeltaLam[nadj*s+i],-h*b[i],ts->vecs_drdy[nadj]);CHKERRQ(ierr);
601       }
602     }
603 
604     /* Stage values of mu */
605     if(ts->vecs_sensip) {
606       ierr = TSAdjointComputeRHSJacobian(ts,rk->stage_time,Y[i],ts->Jacp);CHKERRQ(ierr);
607       if (ts->vec_costintegral) {
608         ierr = TSAdjointComputeDRDPFunction(ts,rk->stage_time,Y[i],ts->vecs_drdp);CHKERRQ(ierr);
609       }
610 
611       for (nadj=0; nadj<ts->numcost; nadj++) {
612         ierr = MatMultTranspose(ts->Jacp,VecSensiTemp[nadj],VecDeltaMu[nadj*s+i]);CHKERRQ(ierr);
613         if (ts->vec_costintegral) {
614           ierr = VecAXPY(VecDeltaMu[nadj*s+i],-h*b[i],ts->vecs_drdp[nadj]);CHKERRQ(ierr);
615         }
616       }
617     }
618   }
619 
620   for (j=0; j<s; j++) w[j] = 1.0;
621   for (nadj=0; nadj<ts->numcost; nadj++) {
622     ierr = VecMAXPY(ts->vecs_sensi[nadj],s,w,&VecDeltaLam[nadj*s]);CHKERRQ(ierr);
623     if(ts->vecs_sensip) {
624       ierr = VecMAXPY(ts->vecs_sensip[nadj],s,w,&VecDeltaMu[nadj*s]);CHKERRQ(ierr);
625     }
626   }
627   rk->status = TS_STEP_COMPLETE;
628   PetscFunctionReturn(0);
629 }
630 
631 static PetscErrorCode TSInterpolate_RK(TS ts,PetscReal itime,Vec X)
632 {
633   TS_RK           *rk = (TS_RK*)ts->data;
634   PetscInt         s  = rk->tableau->s,p = rk->tableau->p,i,j;
635   PetscReal        h;
636   PetscReal        tt,t;
637   PetscScalar     *b;
638   const PetscReal *B = rk->tableau->binterp;
639   PetscErrorCode   ierr;
640 
641   PetscFunctionBegin;
642   if (!B) SETERRQ1(PetscObjectComm((PetscObject)ts),PETSC_ERR_SUP,"TSRK %s does not have an interpolation formula",rk->tableau->name);
643 
644   switch (rk->status) {
645   case TS_STEP_INCOMPLETE:
646   case TS_STEP_PENDING:
647     h = ts->time_step;
648     t = (itime - ts->ptime)/h;
649     break;
650   case TS_STEP_COMPLETE:
651     h = ts->ptime - ts->ptime_prev;
652     t = (itime - ts->ptime)/h + 1; /* In the interval [0,1] */
653     break;
654   default: SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_PLIB,"Invalid TSStepStatus");
655   }
656   ierr = PetscMalloc1(s,&b);CHKERRQ(ierr);
657   for (i=0; i<s; i++) b[i] = 0;
658   for (j=0,tt=t; j<p; j++,tt*=t) {
659     for (i=0; i<s; i++) {
660       b[i]  += h * B[i*p+j] * tt;
661     }
662   }
663 
664   ierr = VecCopy(rk->Y[0],X);CHKERRQ(ierr);
665   ierr = VecMAXPY(X,s,b,rk->YdotRHS);CHKERRQ(ierr);
666 
667   ierr = PetscFree(b);CHKERRQ(ierr);
668   PetscFunctionReturn(0);
669 }
670 
671 /*------------------------------------------------------------*/
672 
673 static PetscErrorCode TSRKTableauReset(TS ts)
674 {
675   TS_RK          *rk = (TS_RK*)ts->data;
676   RKTableau      tab = rk->tableau;
677   PetscErrorCode ierr;
678 
679   PetscFunctionBegin;
680   if (!tab) PetscFunctionReturn(0);
681   ierr = PetscFree(rk->work);CHKERRQ(ierr);
682   ierr = VecDestroyVecs(tab->s,&rk->Y);CHKERRQ(ierr);
683   ierr = VecDestroyVecs(tab->s,&rk->YdotRHS);CHKERRQ(ierr);
684   ierr = VecDestroyVecs(tab->s*ts->numcost,&rk->VecDeltaLam);CHKERRQ(ierr);
685   ierr = VecDestroyVecs(tab->s*ts->numcost,&rk->VecDeltaMu);CHKERRQ(ierr);
686   PetscFunctionReturn(0);
687 }
688 
689 static PetscErrorCode TSReset_RK(TS ts)
690 {
691   TS_RK         *rk = (TS_RK*)ts->data;
692   PetscErrorCode ierr;
693 
694   PetscFunctionBegin;
695   ierr = TSRKTableauReset(ts);CHKERRQ(ierr);
696   ierr = VecDestroy(&rk->VecCostIntegral0);CHKERRQ(ierr);
697   ierr = VecDestroyVecs(ts->numcost,&rk->VecSensiTemp);CHKERRQ(ierr);
698   PetscFunctionReturn(0);
699 }
700 
701 static PetscErrorCode TSDestroy_RK(TS ts)
702 {
703   PetscErrorCode ierr;
704 
705   PetscFunctionBegin;
706   ierr = TSReset_RK(ts);CHKERRQ(ierr);
707   ierr = PetscFree(ts->data);CHKERRQ(ierr);
708   ierr = PetscObjectComposeFunction((PetscObject)ts,"TSRKGetType_C",NULL);CHKERRQ(ierr);
709   ierr = PetscObjectComposeFunction((PetscObject)ts,"TSRKSetType_C",NULL);CHKERRQ(ierr);
710   PetscFunctionReturn(0);
711 }
712 
713 
714 static PetscErrorCode DMCoarsenHook_TSRK(DM fine,DM coarse,void *ctx)
715 {
716   PetscFunctionBegin;
717   PetscFunctionReturn(0);
718 }
719 
720 static PetscErrorCode DMRestrictHook_TSRK(DM fine,Mat restrct,Vec rscale,Mat inject,DM coarse,void *ctx)
721 {
722   PetscFunctionBegin;
723   PetscFunctionReturn(0);
724 }
725 
726 
727 static PetscErrorCode DMSubDomainHook_TSRK(DM dm,DM subdm,void *ctx)
728 {
729   PetscFunctionBegin;
730   PetscFunctionReturn(0);
731 }
732 
733 static PetscErrorCode DMSubDomainRestrictHook_TSRK(DM dm,VecScatter gscat,VecScatter lscat,DM subdm,void *ctx)
734 {
735 
736   PetscFunctionBegin;
737   PetscFunctionReturn(0);
738 }
739 /*
740 static PetscErrorCode RKSetAdjCoe(RKTableau tab)
741 {
742   PetscReal *A,*b;
743   PetscInt        s,i,j;
744   PetscErrorCode  ierr;
745 
746   PetscFunctionBegin;
747   s    = tab->s;
748   ierr = PetscMalloc2(s*s,&A,s,&b);CHKERRQ(ierr);
749 
750   for (i=0; i<s; i++)
751     for (j=0; j<s; j++) {
752       A[i*s+j] = (tab->b[s-1-i]==0) ? 0: -tab->A[s-1-i+(s-1-j)*s] * tab->b[s-1-j] / tab->b[s-1-i];
753       ierr = PetscPrintf(PETSC_COMM_WORLD,"Coefficients: A[%D][%D]=%.6f\n",i,j,A[i*s+j]);CHKERRQ(ierr);
754     }
755 
756   for (i=0; i<s; i++) b[i] = (tab->b[s-1-i]==0)? 0: -tab->b[s-1-i];
757 
758   ierr  = PetscMemcpy(tab->A,A,s*s*sizeof(A[0]));CHKERRQ(ierr);
759   ierr  = PetscMemcpy(tab->b,b,s*sizeof(b[0]));CHKERRQ(ierr);
760   ierr  = PetscFree2(A,b);CHKERRQ(ierr);
761   PetscFunctionReturn(0);
762 }
763 */
764 
765 static PetscErrorCode TSRKTableauSetUp(TS ts)
766 {
767   TS_RK         *rk  = (TS_RK*)ts->data;
768   RKTableau      tab = rk->tableau;
769   PetscErrorCode ierr;
770 
771   PetscFunctionBegin;
772   ierr = PetscMalloc1(tab->s,&rk->work);CHKERRQ(ierr);
773   ierr = VecDuplicateVecs(ts->vec_sol,tab->s,&rk->Y);CHKERRQ(ierr);
774   ierr = VecDuplicateVecs(ts->vec_sol,tab->s,&rk->YdotRHS);CHKERRQ(ierr);
775   PetscFunctionReturn(0);
776 }
777 
778 
779 static PetscErrorCode TSSetUp_RK(TS ts)
780 {
781   TS_RK         *rk = (TS_RK*)ts->data;
782   PetscErrorCode ierr;
783   DM             dm;
784 
785   PetscFunctionBegin;
786   ierr = TSCheckImplicitTerm(ts);CHKERRQ(ierr);
787   ierr = TSRKTableauSetUp(ts);CHKERRQ(ierr);
788   if (!rk->VecCostIntegral0 && ts->vec_costintegral && ts->costintegralfwd) { /* back up cost integral */
789     ierr = VecDuplicate(ts->vec_costintegral,&rk->VecCostIntegral0);CHKERRQ(ierr);
790   }
791   ierr = TSGetDM(ts,&dm);CHKERRQ(ierr);
792   if (dm) {
793     ierr = DMCoarsenHookAdd(dm,DMCoarsenHook_TSRK,DMRestrictHook_TSRK,ts);CHKERRQ(ierr);
794     ierr = DMSubDomainHookAdd(dm,DMSubDomainHook_TSRK,DMSubDomainRestrictHook_TSRK,ts);CHKERRQ(ierr);
795   }
796   PetscFunctionReturn(0);
797 }
798 
799 
800 /*------------------------------------------------------------*/
801 
802 static PetscErrorCode TSSetFromOptions_RK(PetscOptionItems *PetscOptionsObject,TS ts)
803 {
804   TS_RK         *rk = (TS_RK*)ts->data;
805   PetscErrorCode ierr;
806 
807   PetscFunctionBegin;
808   ierr = PetscOptionsHead(PetscOptionsObject,"RK ODE solver options");CHKERRQ(ierr);
809   {
810     RKTableauLink  link;
811     PetscInt       count,choice;
812     PetscBool      flg;
813     const char   **namelist;
814     for (link=RKTableauList,count=0; link; link=link->next,count++) ;
815     ierr = PetscMalloc1(count,&namelist);CHKERRQ(ierr);
816     for (link=RKTableauList,count=0; link; link=link->next,count++) namelist[count] = link->tab.name;
817     ierr = PetscOptionsEList("-ts_rk_type","Family of RK method","TSRKSetType",(const char*const*)namelist,count,rk->tableau->name,&choice,&flg);CHKERRQ(ierr);
818     if (flg) {ierr = TSRKSetType(ts,namelist[choice]);CHKERRQ(ierr);}
819     ierr = PetscFree(namelist);CHKERRQ(ierr);
820   }
821   ierr = PetscOptionsTail();CHKERRQ(ierr);
822   PetscFunctionReturn(0);
823 }
824 
825 static PetscErrorCode PetscFormatRealArray(char buf[],size_t len,const char *fmt,PetscInt n,const PetscReal x[])
826 {
827   PetscErrorCode ierr;
828   PetscInt       i;
829   size_t         left,count;
830   char           *p;
831 
832   PetscFunctionBegin;
833   for (i=0,p=buf,left=len; i<n; i++) {
834     ierr = PetscSNPrintfCount(p,left,fmt,&count,x[i]);CHKERRQ(ierr);
835     if (count >= left) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_OUTOFRANGE,"Insufficient space in buffer");
836     left -= count;
837     p    += count;
838     *p++  = ' ';
839   }
840   p[i ? 0 : -1] = 0;
841   PetscFunctionReturn(0);
842 }
843 
844 static PetscErrorCode TSView_RK(TS ts,PetscViewer viewer)
845 {
846   TS_RK          *rk = (TS_RK*)ts->data;
847   PetscBool      iascii;
848   PetscErrorCode ierr;
849 
850   PetscFunctionBegin;
851   ierr = PetscObjectTypeCompare((PetscObject)viewer,PETSCVIEWERASCII,&iascii);CHKERRQ(ierr);
852   if (iascii) {
853     RKTableau tab  = rk->tableau;
854     TSRKType  rktype;
855     char      buf[512];
856     ierr = TSRKGetType(ts,&rktype);CHKERRQ(ierr);
857     ierr = PetscViewerASCIIPrintf(viewer,"  RK: %s\n",rktype);CHKERRQ(ierr);
858     ierr = PetscViewerASCIIPrintf(viewer,"  Order: %D\n",tab->order);CHKERRQ(ierr);
859     ierr = PetscViewerASCIIPrintf(viewer,"  FSAL property: %s\n",tab->FSAL ? "yes" : "no");CHKERRQ(ierr);
860     ierr = PetscFormatRealArray(buf,sizeof(buf),"% 8.6f",tab->s,tab->c);CHKERRQ(ierr);
861     ierr = PetscViewerASCIIPrintf(viewer,"  Abscissa c = %s\n",buf);CHKERRQ(ierr);
862   }
863   if (ts->adapt) {ierr = TSAdaptView(ts->adapt,viewer);CHKERRQ(ierr);}
864   PetscFunctionReturn(0);
865 }
866 
867 static PetscErrorCode TSLoad_RK(TS ts,PetscViewer viewer)
868 {
869   PetscErrorCode ierr;
870   TSAdapt        adapt;
871 
872   PetscFunctionBegin;
873   ierr = TSGetAdapt(ts,&adapt);CHKERRQ(ierr);
874   ierr = TSAdaptLoad(adapt,viewer);CHKERRQ(ierr);
875   PetscFunctionReturn(0);
876 }
877 
878 /*@C
879   TSRKSetType - Set the type of RK scheme
880 
881   Logically collective
882 
883   Input Parameter:
884 +  ts - timestepping context
885 -  rktype - type of RK-scheme
886 
887   Level: intermediate
888 
889 .seealso: TSRKGetType(), TSRK, TSRK2, TSRK3, TSRKPRSSP2, TSRK4, TSRK5
890 @*/
891 PetscErrorCode TSRKSetType(TS ts,TSRKType rktype)
892 {
893   PetscErrorCode ierr;
894 
895   PetscFunctionBegin;
896   PetscValidHeaderSpecific(ts,TS_CLASSID,1);
897   ierr = PetscTryMethod(ts,"TSRKSetType_C",(TS,TSRKType),(ts,rktype));CHKERRQ(ierr);
898   PetscFunctionReturn(0);
899 }
900 
901 /*@C
902   TSRKGetType - Get the type of RK scheme
903 
904   Logically collective
905 
906   Input Parameter:
907 .  ts - timestepping context
908 
909   Output Parameter:
910 .  rktype - type of RK-scheme
911 
912   Level: intermediate
913 
914 .seealso: TSRKGetType()
915 @*/
916 PetscErrorCode TSRKGetType(TS ts,TSRKType *rktype)
917 {
918   PetscErrorCode ierr;
919 
920   PetscFunctionBegin;
921   PetscValidHeaderSpecific(ts,TS_CLASSID,1);
922   ierr = PetscUseMethod(ts,"TSRKGetType_C",(TS,TSRKType*),(ts,rktype));CHKERRQ(ierr);
923   PetscFunctionReturn(0);
924 }
925 
926 static PetscErrorCode TSRKGetType_RK(TS ts,TSRKType *rktype)
927 {
928   TS_RK *rk = (TS_RK*)ts->data;
929 
930   PetscFunctionBegin;
931   *rktype = rk->tableau->name;
932   PetscFunctionReturn(0);
933 }
934 static PetscErrorCode TSRKSetType_RK(TS ts,TSRKType rktype)
935 {
936   TS_RK          *rk = (TS_RK*)ts->data;
937   PetscErrorCode ierr;
938   PetscBool      match;
939   RKTableauLink  link;
940 
941   PetscFunctionBegin;
942   if (rk->tableau) {
943     ierr = PetscStrcmp(rk->tableau->name,rktype,&match);CHKERRQ(ierr);
944     if (match) PetscFunctionReturn(0);
945   }
946   for (link = RKTableauList; link; link=link->next) {
947     ierr = PetscStrcmp(link->tab.name,rktype,&match);CHKERRQ(ierr);
948     if (match) {
949       if (ts->setupcalled) {ierr = TSRKTableauReset(ts);CHKERRQ(ierr);}
950       rk->tableau = &link->tab;
951       if (ts->setupcalled) {ierr = TSRKTableauSetUp(ts);CHKERRQ(ierr);}
952       PetscFunctionReturn(0);
953     }
954   }
955   SETERRQ1(PetscObjectComm((PetscObject)ts),PETSC_ERR_ARG_UNKNOWN_TYPE,"Could not find '%s'",rktype);
956   PetscFunctionReturn(0);
957 }
958 
959 static PetscErrorCode  TSGetStages_RK(TS ts,PetscInt *ns,Vec **Y)
960 {
961   TS_RK          *rk = (TS_RK*)ts->data;
962 
963   PetscFunctionBegin;
964   *ns = rk->tableau->s;
965   if(Y) *Y  = rk->Y;
966   PetscFunctionReturn(0);
967 }
968 
969 
970 /* ------------------------------------------------------------ */
971 /*MC
972       TSRK - ODE and DAE solver using Runge-Kutta schemes
973 
974   The user should provide the right hand side of the equation
975   using TSSetRHSFunction().
976 
977   Notes:
978   The default is TSRK3, it can be changed with TSRKSetType() or -ts_rk_type
979 
980   Level: beginner
981 
982 .seealso:  TSCreate(), TS, TSSetType(), TSRKSetType(), TSRKGetType(), TSRKSetFullyImplicit(), TSRK2D, TTSRK2E, TSRK3,
983            TSRK4, TSRK5, TSRKPRSSP2, TSRKBPR3, TSRKType, TSRKRegister()
984 
985 M*/
986 PETSC_EXTERN PetscErrorCode TSCreate_RK(TS ts)
987 {
988   TS_RK          *rk;
989   PetscErrorCode ierr;
990 
991   PetscFunctionBegin;
992   ierr = TSRKInitializePackage();CHKERRQ(ierr);
993 
994   ts->ops->reset          = TSReset_RK;
995   ts->ops->destroy        = TSDestroy_RK;
996   ts->ops->view           = TSView_RK;
997   ts->ops->load           = TSLoad_RK;
998   ts->ops->setup          = TSSetUp_RK;
999   ts->ops->adjointsetup   = TSAdjointSetUp_RK;
1000   ts->ops->step           = TSStep_RK;
1001   ts->ops->interpolate    = TSInterpolate_RK;
1002   ts->ops->evaluatestep   = TSEvaluateStep_RK;
1003   ts->ops->rollback       = TSRollBack_RK;
1004   ts->ops->setfromoptions = TSSetFromOptions_RK;
1005   ts->ops->getstages      = TSGetStages_RK;
1006   ts->ops->adjointstep    = TSAdjointStep_RK;
1007 
1008   ts->ops->adjointintegral = TSAdjointCostIntegral_RK;
1009   ts->ops->forwardintegral = TSForwardCostIntegral_RK;
1010 
1011   ierr = PetscNewLog(ts,&rk);CHKERRQ(ierr);
1012   ts->data = (void*)rk;
1013 
1014   ierr = PetscObjectComposeFunction((PetscObject)ts,"TSRKGetType_C",TSRKGetType_RK);CHKERRQ(ierr);
1015   ierr = PetscObjectComposeFunction((PetscObject)ts,"TSRKSetType_C",TSRKSetType_RK);CHKERRQ(ierr);
1016 
1017   ierr = TSRKSetType(ts,TSRKDefault);CHKERRQ(ierr);
1018   PetscFunctionReturn(0);
1019 }
1020