xref: /petsc/src/ts/impls/arkimex/arkimex.c (revision 8cd211a4b3bd7a7e8a5ce60d0629aecc5bd5c55c)
1 /*
2   Code for timestepping with additive Runge-Kutta IMEX method
3 
4   Notes:
5   The general system is written as
6 
7   F(t,X,Xdot) = G(t,X)
8 
9   where F represents the stiff part of the physics and G represents the non-stiff part.
10 
11 */
12 #include <private/tsimpl.h>                /*I   "petscts.h"   I*/
13 
14 static const TSARKIMEXType TSARKIMEXDefault = TSARKIMEX2E;
15 static PetscBool TSARKIMEXRegisterAllCalled;
16 static PetscBool TSARKIMEXPackageInitialized;
17 
18 typedef struct _ARKTableau *ARKTableau;
19 struct _ARKTableau {
20   char *name;
21   PetscInt order;               /* Classical approximation order of the method */
22   PetscInt s;                   /* Number of stages */
23   PetscInt pinterp;             /* Interpolation order */
24   PetscReal *At,*bt,*ct;        /* Stiff tableau */
25   PetscReal *A,*b,*c;           /* Non-stiff tableau */
26   PetscReal *binterpt,*binterp; /* Dense output formula */
27 };
28 typedef struct _ARKTableauLink *ARKTableauLink;
29 struct _ARKTableauLink {
30   struct _ARKTableau tab;
31   ARKTableauLink next;
32 };
33 static ARKTableauLink ARKTableauList;
34 
35 typedef struct {
36   ARKTableau  tableau;
37   Vec         *Y;               /* States computed during the step */
38   Vec         *YdotI;           /* Time derivatives for the stiff part */
39   Vec         *YdotRHS;         /* Function evaluations for the non-stiff part */
40   Vec         Ydot;             /* Work vector holding Ydot during residual evaluation */
41   Vec         Work;             /* Generic work vector */
42   Vec         Z;                /* Ydot = shift(Y-Z) */
43   PetscScalar *work;            /* Scalar work */
44   PetscReal   shift;
45   PetscReal   stage_time;
46   PetscBool   imex;
47 } TS_ARKIMEX;
48 
49 /*MC
50      TSARKIMEX2D - Second order ARK IMEX scheme with L-stable implicit part.
51 
52      This method has one explicit stage and two implicit stages.
53 
54      Level: advanced
55 
56 .seealso: TSARKIMEX
57 M*/
58 /*MC
59      TSARKIMEX2E - Second order ARK IMEX scheme with L-stable implicit part.
60 
61      This method has one explicit stage and two implicit stages. It is is an optimal method developed by Emil Constantinescu.
62 
63      Level: advanced
64 
65 .seealso: TSARKIMEX
66 M*/
67 /*MC
68      TSARKIMEX3 - Third order ARK IMEX scheme with L-stable implicit part.
69 
70      This method has one explicit stage and three implicit stages.
71 
72      References:
73      Kennedy and Carpenter 2003.
74 
75      Level: advanced
76 
77 .seealso: TSARKIMEX
78 M*/
79 /*MC
80      TSARKIMEX4 - Fourth order ARK IMEX scheme with L-stable implicit part.
81 
82      This method has one explicit stage and four implicit stages.
83 
84      References:
85      Kennedy and Carpenter 2003.
86 
87      Level: advanced
88 
89 .seealso: TSARKIMEX
90 M*/
91 /*MC
92      TSARKIMEX5 - Fifth order ARK IMEX scheme with L-stable implicit part.
93 
94      This method has one explicit stage and five implicit stages.
95 
96      References:
97      Kennedy and Carpenter 2003.
98 
99      Level: advanced
100 
101 .seealso: TSARKIMEX
102 M*/
103 
104 #undef __FUNCT__
105 #define __FUNCT__ "TSARKIMEXRegisterAll"
106 /*@C
107   TSARKIMEXRegisterAll - Registers all of the additive Runge-Kutta implicit-explicit methods in TSARKIMEX
108 
109   Not Collective, but should be called by all processes which will need the schemes to be registered
110 
111   Level: advanced
112 
113 .keywords: TS, TSARKIMEX, register, all
114 
115 .seealso:  TSARKIMEXRegisterDestroy()
116 @*/
117 PetscErrorCode TSARKIMEXRegisterAll(void)
118 {
119   PetscErrorCode ierr;
120 
121   PetscFunctionBegin;
122   if (TSARKIMEXRegisterAllCalled) PetscFunctionReturn(0);
123   TSARKIMEXRegisterAllCalled = PETSC_TRUE;
124 
125   {
126     const PetscReal
127       A[3][3] = {{0,0,0},
128                  {0.41421356237309504880,0,0},
129                  {0.75,0.25,0}},
130       At[3][3] = {{0,0,0},
131                   {0.12132034355964257320,0.29289321881345247560,0},
132                   {0.20710678118654752440,0.50000000000000000000,0.29289321881345247560}},
133       binterpt[3][2] = {{1,-0.5},{0,0},{0,0.5}};
134     ierr = TSARKIMEXRegister(TSARKIMEX2D,2,3,&At[0][0],PETSC_NULL,PETSC_NULL,&A[0][0],PETSC_NULL,PETSC_NULL,2,binterpt[0],PETSC_NULL);CHKERRQ(ierr);
135   }
136   {                             /* Optimal for linear implicit part */
137     const PetscReal s2 = PetscSqrtReal((PetscReal)2.0),
138       A[3][3] = {{0,0,0},
139                  {2-s2,0,0},
140                  {(3-2*s2)/6,(3+2*s2)/6,0}},
141       At[3][3] = {{0,0,0},
142                   {1-1/s2,1-1/s2,0},
143                   {1/(2*s2),1/(2*s2),1-1/s2}},
144       binterpt[3][2] = {{1,-0.5},{0,0},{0,0.5}};
145     ierr = TSARKIMEXRegister(TSARKIMEX2E,2,3,&At[0][0],PETSC_NULL,PETSC_NULL,&A[0][0],PETSC_NULL,PETSC_NULL,2,binterpt[0],PETSC_NULL);CHKERRQ(ierr);
146   }
147   {
148     const PetscReal
149       A[4][4] = {{0,0,0,0},
150                  {1767732205903./2027836641118.,0,0,0},
151                  {5535828885825./10492691773637.,788022342437./10882634858940.,0,0},
152                  {6485989280629./16251701735622.,-4246266847089./9704473918619.,10755448449292./10357097424841.,0}},
153       At[4][4] = {{0,0,0,0},
154                   {1767732205903./4055673282236.,1767732205903./4055673282236.,0,0},
155                   {2746238789719./10658868560708.,-640167445237./6845629431997.,1767732205903./4055673282236.,0},
156                   {1471266399579./7840856788654.,-4482444167858./7529755066697.,11266239266428./11593286722821.,1767732205903./4055673282236.}},
157       binterpt[4][2] = {{4655552711362./22874653954995., -215264564351./13552729205753.},
158                         {-18682724506714./9892148508045.,17870216137069./13817060693119.},
159                         {34259539580243./13192909600954.,-28141676662227./17317692491321.},
160                         {584795268549./6622622206610.,   2508943948391./7218656332882.}};
161     ierr = TSARKIMEXRegister(TSARKIMEX3,3,4,&At[0][0],PETSC_NULL,PETSC_NULL,&A[0][0],PETSC_NULL,PETSC_NULL,2,binterpt[0],PETSC_NULL);CHKERRQ(ierr);
162   }
163   {
164     const PetscReal
165       A[6][6] = {{0,0,0,0,0,0},
166                  {1./2,0,0,0,0,0},
167                  {13861./62500.,6889./62500.,0,0,0,0},
168                  {-116923316275./2393684061468.,-2731218467317./15368042101831.,9408046702089./11113171139209.,0,0,0},
169                  {-451086348788./2902428689909.,-2682348792572./7519795681897.,12662868775082./11960479115383.,3355817975965./11060851509271.,0,0},
170                  {647845179188./3216320057751.,73281519250./8382639484533.,552539513391./3454668386233.,3354512671639./8306763924573.,4040./17871.,0}},
171       At[6][6] = {{0,0,0,0,0,0},
172                   {1./4,1./4,0,0,0,0},
173                   {8611./62500.,-1743./31250.,1./4,0,0,0},
174                   {5012029./34652500.,-654441./2922500.,174375./388108.,1./4,0,0},
175                   {15267082809./155376265600.,-71443401./120774400.,730878875./902184768.,2285395./8070912.,1./4,0},
176                   {82889./524892.,0,15625./83664.,69875./102672.,-2260./8211,1./4}},
177       binterpt[6][3] = {{6943876665148./7220017795957.,-54480133./30881146.,6818779379841./7100303317025.},
178                         {0,0,0},
179                         {7640104374378./9702883013639.,-11436875./14766696.,2173542590792./12501825683035.},
180                         {-20649996744609./7521556579894.,174696575./18121608.,-31592104683404./5083833661969.},
181                         {8854892464581./2390941311638.,-12120380./966161.,61146701046299./7138195549469.},
182                         {-11397109935349./6675773540249.,3843./706.,-17219254887155./4939391667607.}};
183     ierr = TSARKIMEXRegister(TSARKIMEX4,4,6,&At[0][0],PETSC_NULL,PETSC_NULL,&A[0][0],PETSC_NULL,PETSC_NULL,3,binterpt[0],PETSC_NULL);CHKERRQ(ierr);
184   }
185   {
186     const PetscReal
187       A[8][8] = {{0,0,0,0,0,0,0,0},
188                  {41./100,0,0,0,0,0,0,0},
189                  {367902744464./2072280473677.,677623207551./8224143866563.,0,0,0,0,0,0},
190                  {1268023523408./10340822734521.,0,1029933939417./13636558850479.,0,0,0,0,0},
191                  {14463281900351./6315353703477.,0,66114435211212./5879490589093.,-54053170152839./4284798021562.,0,0,0,0},
192                  {14090043504691./34967701212078.,0,15191511035443./11219624916014.,-18461159152457./12425892160975.,-281667163811./9011619295870.,0,0,0},
193                  {19230459214898./13134317526959.,0,21275331358303./2942455364971.,-38145345988419./4862620318723.,-1./8,-1./8,0,0},
194                  {-19977161125411./11928030595625.,0,-40795976796054./6384907823539.,177454434618887./12078138498510.,782672205425./8267701900261.,-69563011059811./9646580694205.,7356628210526./4942186776405.,0}},
195       At[8][8] = {{0,0,0,0,0,0,0,0},
196                   {41./200.,41./200.,0,0,0,0,0,0},
197                   {41./400.,-567603406766./11931857230679.,41./200.,0,0,0,0,0},
198                   {683785636431./9252920307686.,0,-110385047103./1367015193373.,41./200.,0,0,0,0},
199                   {3016520224154./10081342136671.,0,30586259806659./12414158314087.,-22760509404356./11113319521817.,41./200.,0,0,0},
200                   {218866479029./1489978393911.,0,638256894668./5436446318841.,-1179710474555./5321154724896.,-60928119172./8023461067671.,41./200.,0,0},
201                   {1020004230633./5715676835656.,0,25762820946817./25263940353407.,-2161375909145./9755907335909.,-211217309593./5846859502534.,-4269925059573./7827059040749.,41./200,0},
202                   {-872700587467./9133579230613.,0,0,22348218063261./9555858737531.,-1143369518992./8141816002931.,-39379526789629./19018526304540.,32727382324388./42900044865799.,41./200.}},
203       binterpt[8][3] = {{-17674230611817./10670229744614. ,  43486358583215./12773830924787. , -9257016797708./5021505065439.},
204                         {0                               ,  0                              , 0                            },
205                         {0                               ,  0                              , 0                            },
206                         {65168852399939./7868540260826.  ,  -91478233927265./11067650958493., 26096422576131./11239449250142.},
207                         {15494834004392./5936557850923.  ,  -79368583304911./10890268929626., 92396832856987./20362823103730.},
208                         {-99329723586156./26959484932159.,  -12239297817655./9152339842473. , 30029262896817./10175596800299.},
209                         {-19024464361622./5461577185407. ,  115839755401235./10719374521269., -26136350496073./3983972220547.},
210                         {-6511271360970./6095937251113.  ,  5843115559534./2180450260947.   , -5289405421727./3760307252460. }};
211     ierr = TSARKIMEXRegister(TSARKIMEX5,5,8,&At[0][0],PETSC_NULL,PETSC_NULL,&A[0][0],PETSC_NULL,PETSC_NULL,3,binterpt[0],PETSC_NULL);CHKERRQ(ierr);
212   }
213 
214   PetscFunctionReturn(0);
215 }
216 
217 #undef __FUNCT__
218 #define __FUNCT__ "TSARKIMEXRegisterDestroy"
219 /*@C
220    TSARKIMEXRegisterDestroy - Frees the list of schemes that were registered by TSARKIMEXRegister().
221 
222    Not Collective
223 
224    Level: advanced
225 
226 .keywords: TSARKIMEX, register, destroy
227 .seealso: TSARKIMEXRegister(), TSARKIMEXRegisterAll(), TSARKIMEXRegisterDynamic()
228 @*/
229 PetscErrorCode TSARKIMEXRegisterDestroy(void)
230 {
231   PetscErrorCode ierr;
232   ARKTableauLink link;
233 
234   PetscFunctionBegin;
235   while ((link = ARKTableauList)) {
236     ARKTableau t = &link->tab;
237     ARKTableauList = link->next;
238     ierr = PetscFree6(t->At,t->bt,t->ct,t->A,t->b,t->c);CHKERRQ(ierr);
239     ierr = PetscFree2(t->binterpt,t->binterp);CHKERRQ(ierr);
240     ierr = PetscFree(t->name);CHKERRQ(ierr);
241     ierr = PetscFree(link);CHKERRQ(ierr);
242   }
243   TSARKIMEXRegisterAllCalled = PETSC_FALSE;
244   PetscFunctionReturn(0);
245 }
246 
247 #undef __FUNCT__
248 #define __FUNCT__ "TSARKIMEXInitializePackage"
249 /*@C
250   TSARKIMEXInitializePackage - This function initializes everything in the TSARKIMEX package. It is called
251   from PetscDLLibraryRegister() when using dynamic libraries, and on the first call to TSCreate_ARKIMEX()
252   when using static libraries.
253 
254   Input Parameter:
255   path - The dynamic library path, or PETSC_NULL
256 
257   Level: developer
258 
259 .keywords: TS, TSARKIMEX, initialize, package
260 .seealso: PetscInitialize()
261 @*/
262 PetscErrorCode TSARKIMEXInitializePackage(const char path[])
263 {
264   PetscErrorCode ierr;
265 
266   PetscFunctionBegin;
267   if (TSARKIMEXPackageInitialized) PetscFunctionReturn(0);
268   TSARKIMEXPackageInitialized = PETSC_TRUE;
269   ierr = TSARKIMEXRegisterAll();CHKERRQ(ierr);
270   ierr = PetscRegisterFinalize(TSARKIMEXFinalizePackage);CHKERRQ(ierr);
271   PetscFunctionReturn(0);
272 }
273 
274 #undef __FUNCT__
275 #define __FUNCT__ "TSARKIMEXFinalizePackage"
276 /*@C
277   TSARKIMEXFinalizePackage - This function destroys everything in the TSARKIMEX package. It is
278   called from PetscFinalize().
279 
280   Level: developer
281 
282 .keywords: Petsc, destroy, package
283 .seealso: PetscFinalize()
284 @*/
285 PetscErrorCode TSARKIMEXFinalizePackage(void)
286 {
287   PetscErrorCode ierr;
288 
289   PetscFunctionBegin;
290   TSARKIMEXPackageInitialized = PETSC_FALSE;
291   ierr = TSARKIMEXRegisterDestroy();CHKERRQ(ierr);
292   PetscFunctionReturn(0);
293 }
294 
295 #undef __FUNCT__
296 #define __FUNCT__ "TSARKIMEXRegister"
297 /*@C
298    TSARKIMEXRegister - register an ARK IMEX scheme by providing the entries in the Butcher tableau and optionally embedded approximations and interpolation
299 
300    Not Collective, but the same schemes should be registered on all processes on which they will be used
301 
302    Input Parameters:
303 +  name - identifier for method
304 .  order - approximation order of method
305 .  s - number of stages, this is the dimension of the matrices below
306 .  At - Butcher table of stage coefficients for stiff part (dimension s*s, row-major)
307 .  bt - Butcher table for completing the stiff part of the step (dimension s; PETSC_NULL to use the last row of At)
308 .  ct - Abscissa of each stiff stage (dimension s, PETSC_NULL to use row sums of At)
309 .  A - Non-stiff stage coefficients (dimension s*s, row-major)
310 .  b - Non-stiff step completion table (dimension s; PETSC_NULL to use last row of At)
311 .  c - Non-stiff abscissa (dimension s; PETSC_NULL to use row sums of A)
312 .  pinterp - Order of the interpolation scheme, equal to the number of columns of binterpt and binterp
313 .  binterpt - Coefficients of the interpolation formula for the stiff part (dimension s*pinterp)
314 -  binterp - Coefficients of the interpolation formula for the non-stiff part (dimension s*pinterp; PETSC_NULL to reuse binterpt)
315 
316    Notes:
317    Several ARK IMEX methods are provided, this function is only needed to create new methods.
318 
319    Level: advanced
320 
321 .keywords: TS, register
322 
323 .seealso: TSARKIMEX
324 @*/
325 PetscErrorCode TSARKIMEXRegister(const TSARKIMEXType name,PetscInt order,PetscInt s,
326                                  const PetscReal At[],const PetscReal bt[],const PetscReal ct[],
327                                  const PetscReal A[],const PetscReal b[],const PetscReal c[],
328                                  PetscInt pinterp,const PetscReal binterpt[],const PetscReal binterp[])
329 {
330   PetscErrorCode ierr;
331   ARKTableauLink link;
332   ARKTableau t;
333   PetscInt i,j;
334 
335   PetscFunctionBegin;
336   ierr = PetscMalloc(sizeof(*link),&link);CHKERRQ(ierr);
337   ierr = PetscMemzero(link,sizeof(*link));CHKERRQ(ierr);
338   t = &link->tab;
339   ierr = PetscStrallocpy(name,&t->name);CHKERRQ(ierr);
340   t->order = order;
341   t->s = s;
342   ierr = PetscMalloc6(s*s,PetscReal,&t->At,s,PetscReal,&t->bt,s,PetscReal,&t->ct,s*s,PetscReal,&t->A,s,PetscReal,&t->b,s,PetscReal,&t->c);CHKERRQ(ierr);
343   ierr = PetscMemcpy(t->At,At,s*s*sizeof(At[0]));CHKERRQ(ierr);
344   ierr = PetscMemcpy(t->A,A,s*s*sizeof(A[0]));CHKERRQ(ierr);
345   if (bt) {ierr = PetscMemcpy(t->bt,bt,s*sizeof(bt[0]));CHKERRQ(ierr);}
346   else for (i=0; i<s; i++) t->bt[i] = At[(s-1)*s+i];
347   if (b) {ierr = PetscMemcpy(t->b,b,s*sizeof(b[0]));CHKERRQ(ierr);}
348   else for (i=0; i<s; i++) t->b[i] = At[(s-1)*s+i];
349   if (ct) {ierr = PetscMemcpy(t->ct,ct,s*sizeof(ct[0]));CHKERRQ(ierr);}
350   else for (i=0; i<s; i++) for (j=0,t->ct[i]=0; j<s; j++) t->ct[i] += At[i*s+j];
351   if (c) {ierr = PetscMemcpy(t->c,c,s*sizeof(c[0]));CHKERRQ(ierr);}
352   else for (i=0; i<s; i++) for (j=0,t->c[i]=0; j<s; j++) t->c[i] += A[i*s+j];
353   t->pinterp = pinterp;
354   ierr = PetscMalloc2(s*pinterp,PetscReal,&t->binterpt,s*pinterp,PetscReal,&t->binterp);CHKERRQ(ierr);
355   ierr = PetscMemcpy(t->binterpt,binterpt,s*pinterp*sizeof(binterpt[0]));CHKERRQ(ierr);
356   ierr = PetscMemcpy(t->binterp,binterp?binterp:binterpt,s*pinterp*sizeof(binterpt[0]));CHKERRQ(ierr);
357   link->next = ARKTableauList;
358   ARKTableauList = link;
359   PetscFunctionReturn(0);
360 }
361 
362 #undef __FUNCT__
363 #define __FUNCT__ "TSStep_ARKIMEX"
364 static PetscErrorCode TSStep_ARKIMEX(TS ts)
365 {
366   TS_ARKIMEX      *ark = (TS_ARKIMEX*)ts->data;
367   ARKTableau      tab  = ark->tableau;
368   const PetscInt  s    = tab->s;
369   const PetscReal *At  = tab->At,*A = tab->A,*bt = tab->bt,*b = tab->b,*ct = tab->ct,*c = tab->c;
370   PetscScalar     *w   = ark->work;
371   Vec             *Y   = ark->Y,*YdotI = ark->YdotI,*YdotRHS = ark->YdotRHS,Ydot = ark->Ydot,W = ark->Work,Z = ark->Z;
372   SNES            snes;
373   PetscInt        i,j,its,lits;
374   PetscReal       next_time_step;
375   PetscReal       h,t;
376   PetscErrorCode  ierr;
377 
378   PetscFunctionBegin;
379   ierr = TSGetSNES(ts,&snes);CHKERRQ(ierr);
380   next_time_step = ts->time_step;
381   h = ts->time_step;
382   t = ts->ptime;
383 
384   for (i=0; i<s; i++) {
385     if (At[i*s+i] == 0) {           /* This stage is explicit */
386       ierr = VecCopy(ts->vec_sol,Y[i]);CHKERRQ(ierr);
387       for (j=0; j<i; j++) w[j] = -h*At[i*s+j];
388       ierr = VecMAXPY(Y[i],i,w,YdotI);CHKERRQ(ierr);
389       for (j=0; j<i; j++) w[j] = h*A[i*s+j];
390       ierr = VecMAXPY(Y[i],i,w,YdotRHS);CHKERRQ(ierr);
391     } else {
392       ark->stage_time = t + h*ct[i];
393       ark->shift = 1./(h*At[i*s+i]);
394       /* Affine part */
395       ierr = VecZeroEntries(W);CHKERRQ(ierr);
396       for (j=0; j<i; j++) w[j] = h*A[i*s+j];
397       ierr = VecMAXPY(W,i,w,YdotRHS);CHKERRQ(ierr);
398       /* Ydot = shift*(Y-Z) */
399       ierr = VecCopy(ts->vec_sol,Z);CHKERRQ(ierr);
400       for (j=0; j<i; j++) w[j] = -h*At[i*s+j];
401       ierr = VecMAXPY(Z,i,w,YdotI);CHKERRQ(ierr);
402       /* Initial guess taken from last stage */
403       ierr = VecCopy(i>0?Y[i-1]:ts->vec_sol,Y[i]);CHKERRQ(ierr);
404       ierr = SNESSolve(snes,W,Y[i]);CHKERRQ(ierr);
405       ierr = SNESGetIterationNumber(snes,&its);CHKERRQ(ierr);
406       ierr = SNESGetLinearSolveIterations(snes,&lits);CHKERRQ(ierr);
407       ts->nonlinear_its += its; ts->linear_its += lits;
408     }
409     ierr = VecZeroEntries(Ydot);CHKERRQ(ierr);
410     ierr = TSComputeIFunction(ts,t+h*ct[i],Y[i],Ydot,YdotI[i],ark->imex);CHKERRQ(ierr);
411     if (ark->imex) {
412       ierr = TSComputeRHSFunction(ts,t+h*c[i],Y[i],YdotRHS[i]);CHKERRQ(ierr);
413     } else {
414       ierr = VecZeroEntries(YdotRHS[i]);CHKERRQ(ierr);
415     }
416   }
417   for (j=0; j<s; j++) w[j] = -h*bt[j];
418   ierr = VecMAXPY(ts->vec_sol,s,w,YdotI);CHKERRQ(ierr);
419   for (j=0; j<s; j++) w[j] = h*b[j];
420   ierr = VecMAXPY(ts->vec_sol,s,w,YdotRHS);CHKERRQ(ierr);
421 
422   ts->ptime += ts->time_step;
423   ts->time_step = next_time_step;
424   ts->steps++;
425   PetscFunctionReturn(0);
426 }
427 
428 #undef __FUNCT__
429 #define __FUNCT__ "TSInterpolate_ARKIMEX"
430 static PetscErrorCode TSInterpolate_ARKIMEX(TS ts,PetscReal itime,Vec X)
431 {
432   TS_ARKIMEX *ark = (TS_ARKIMEX*)ts->data;
433   PetscInt s = ark->tableau->s,pinterp = ark->tableau->pinterp,i,j;
434   PetscReal tt,t = (itime - ts->ptime)/ts->time_step + 1; /* In the interval [0,1] */
435   PetscScalar *bt,*b;
436   const PetscReal *Bt = ark->tableau->binterpt,*B = ark->tableau->binterp;
437   PetscErrorCode ierr;
438 
439   PetscFunctionBegin;
440   if (!Bt || !B) SETERRQ1(((PetscObject)ts)->comm,PETSC_ERR_SUP,"TSARKIMEX %s does not have an interpolation formula",ark->tableau->name);
441   ierr = PetscMalloc2(s,PetscScalar,&bt,s,PetscScalar,&b);CHKERRQ(ierr);
442   for (i=0; i<s; i++) bt[i] = b[i] = 0;
443   for (j=0,tt=t; j<pinterp; j++,tt*=t) {
444     for (i=0; i<s; i++) {
445       bt[i] += ts->time_step * Bt[i*pinterp+j] * tt * -1.0;
446       b[i]  += ts->time_step * B[i*pinterp+j] * tt;
447     }
448   }
449   if (ark->tableau->At[0*s+0] != 0.0) SETERRQ(((PetscObject)ts)->comm,PETSC_ERR_SUP,"First stage not explicit so starting stage not saved");
450   ierr = VecCopy(ark->Y[0],X);CHKERRQ(ierr);
451   ierr = VecMAXPY(X,s,bt,ark->YdotI);CHKERRQ(ierr);
452   ierr = VecMAXPY(X,s,b,ark->YdotRHS);CHKERRQ(ierr);
453   ierr = PetscFree2(bt,b);CHKERRQ(ierr);
454   PetscFunctionReturn(0);
455 }
456 
457 /*------------------------------------------------------------*/
458 #undef __FUNCT__
459 #define __FUNCT__ "TSReset_ARKIMEX"
460 static PetscErrorCode TSReset_ARKIMEX(TS ts)
461 {
462   TS_ARKIMEX      *ark = (TS_ARKIMEX*)ts->data;
463   PetscInt        s;
464   PetscErrorCode  ierr;
465 
466   PetscFunctionBegin;
467   if (!ark->tableau) PetscFunctionReturn(0);
468    s = ark->tableau->s;
469   ierr = VecDestroyVecs(s,&ark->Y);CHKERRQ(ierr);
470   ierr = VecDestroyVecs(s,&ark->YdotI);CHKERRQ(ierr);
471   ierr = VecDestroyVecs(s,&ark->YdotRHS);CHKERRQ(ierr);
472   ierr = VecDestroy(&ark->Ydot);CHKERRQ(ierr);
473   ierr = VecDestroy(&ark->Work);CHKERRQ(ierr);
474   ierr = VecDestroy(&ark->Z);CHKERRQ(ierr);
475   ierr = PetscFree(ark->work);CHKERRQ(ierr);
476   PetscFunctionReturn(0);
477 }
478 
479 #undef __FUNCT__
480 #define __FUNCT__ "TSDestroy_ARKIMEX"
481 static PetscErrorCode TSDestroy_ARKIMEX(TS ts)
482 {
483   PetscErrorCode  ierr;
484 
485   PetscFunctionBegin;
486   ierr = TSReset_ARKIMEX(ts);CHKERRQ(ierr);
487   ierr = PetscFree(ts->data);CHKERRQ(ierr);
488   ierr = PetscObjectComposeFunctionDynamic((PetscObject)ts,"TSARKIMEXGetType_C","",PETSC_NULL);CHKERRQ(ierr);
489   ierr = PetscObjectComposeFunctionDynamic((PetscObject)ts,"TSARKIMEXSetType_C","",PETSC_NULL);CHKERRQ(ierr);
490   ierr = PetscObjectComposeFunctionDynamic((PetscObject)ts,"TSARKIMEXSetFullyImplicit_C","",PETSC_NULL);CHKERRQ(ierr);
491   PetscFunctionReturn(0);
492 }
493 
494 /*
495   This defines the nonlinear equation that is to be solved with SNES
496   G(U) = F[t0+Theta*dt, U, (U-U0)*shift] = 0
497 */
498 #undef __FUNCT__
499 #define __FUNCT__ "SNESTSFormFunction_ARKIMEX"
500 static PetscErrorCode SNESTSFormFunction_ARKIMEX(SNES snes,Vec X,Vec F,TS ts)
501 {
502   TS_ARKIMEX     *ark = (TS_ARKIMEX*)ts->data;
503   PetscErrorCode ierr;
504 
505   PetscFunctionBegin;
506   ierr = VecAXPBYPCZ(ark->Ydot,-ark->shift,ark->shift,0,ark->Z,X);CHKERRQ(ierr); /* Ydot = shift*(X-Z) */
507   ierr = TSComputeIFunction(ts,ark->stage_time,X,ark->Ydot,F,ark->imex);CHKERRQ(ierr);
508   PetscFunctionReturn(0);
509 }
510 
511 #undef __FUNCT__
512 #define __FUNCT__ "SNESTSFormJacobian_ARKIMEX"
513 static PetscErrorCode SNESTSFormJacobian_ARKIMEX(SNES snes,Vec X,Mat *A,Mat *B,MatStructure *str,TS ts)
514 {
515   TS_ARKIMEX       *ark = (TS_ARKIMEX*)ts->data;
516   PetscErrorCode ierr;
517 
518   PetscFunctionBegin;
519   /* ark->Ydot has already been computed in SNESTSFormFunction_ARKIMEX (SNES guarantees this) */
520   ierr = TSComputeIJacobian(ts,ark->stage_time,X,ark->Ydot,ark->shift,A,B,str,PETSC_TRUE);CHKERRQ(ierr);
521   PetscFunctionReturn(0);
522 }
523 
524 #undef __FUNCT__
525 #define __FUNCT__ "TSSetUp_ARKIMEX"
526 static PetscErrorCode TSSetUp_ARKIMEX(TS ts)
527 {
528   TS_ARKIMEX     *ark = (TS_ARKIMEX*)ts->data;
529   ARKTableau     tab  = ark->tableau;
530   PetscInt       s = tab->s;
531   PetscErrorCode ierr;
532 
533   PetscFunctionBegin;
534   if (!ark->tableau) {
535     ierr = TSARKIMEXSetType(ts,TSARKIMEXDefault);CHKERRQ(ierr);
536   }
537   ierr = VecDuplicateVecs(ts->vec_sol,s,&ark->Y);CHKERRQ(ierr);
538   ierr = VecDuplicateVecs(ts->vec_sol,s,&ark->YdotI);CHKERRQ(ierr);
539   ierr = VecDuplicateVecs(ts->vec_sol,s,&ark->YdotRHS);CHKERRQ(ierr);
540   ierr = VecDuplicate(ts->vec_sol,&ark->Ydot);CHKERRQ(ierr);
541   ierr = VecDuplicate(ts->vec_sol,&ark->Work);CHKERRQ(ierr);
542   ierr = VecDuplicate(ts->vec_sol,&ark->Z);CHKERRQ(ierr);
543   ierr = PetscMalloc(s*sizeof(ark->work[0]),&ark->work);CHKERRQ(ierr);
544   PetscFunctionReturn(0);
545 }
546 /*------------------------------------------------------------*/
547 
548 #undef __FUNCT__
549 #define __FUNCT__ "TSSetFromOptions_ARKIMEX"
550 static PetscErrorCode TSSetFromOptions_ARKIMEX(TS ts)
551 {
552   TS_ARKIMEX     *ark = (TS_ARKIMEX*)ts->data;
553   PetscErrorCode ierr;
554   char           arktype[256];
555 
556   PetscFunctionBegin;
557   ierr = PetscOptionsHead("ARKIMEX ODE solver options");CHKERRQ(ierr);
558   {
559     ARKTableauLink link;
560     PetscInt count,choice;
561     PetscBool flg;
562     const char **namelist;
563     ierr = PetscStrncpy(arktype,TSARKIMEXDefault,sizeof arktype);CHKERRQ(ierr);
564     for (link=ARKTableauList,count=0; link; link=link->next,count++) ;
565     ierr = PetscMalloc(count*sizeof(char*),&namelist);CHKERRQ(ierr);
566     for (link=ARKTableauList,count=0; link; link=link->next,count++) namelist[count] = link->tab.name;
567     ierr = PetscOptionsEList("-ts_arkimex_type","Family of ARK IMEX method","TSARKIMEXSetType",(const char*const*)namelist,count,arktype,&choice,&flg);CHKERRQ(ierr);
568     ierr = TSARKIMEXSetType(ts,flg ? namelist[choice] : arktype);CHKERRQ(ierr);
569     ierr = PetscFree(namelist);CHKERRQ(ierr);
570     flg = (PetscBool)!ark->imex;
571     ierr = PetscOptionsBool("-ts_arkimex_fully_implicit","Solve the problem fully implicitly","TSARKIMEXSetFullyImplicit",flg,&flg,PETSC_NULL);CHKERRQ(ierr);
572     ark->imex = (PetscBool)!flg;
573     ierr = SNESSetFromOptions(ts->snes);CHKERRQ(ierr);
574   }
575   ierr = PetscOptionsTail();CHKERRQ(ierr);
576   PetscFunctionReturn(0);
577 }
578 
579 #undef __FUNCT__
580 #define __FUNCT__ "PetscFormatRealArray"
581 static PetscErrorCode PetscFormatRealArray(char buf[],size_t len,const char *fmt,PetscInt n,const PetscReal x[])
582 {
583   PetscErrorCode ierr;
584   PetscInt i;
585   size_t left,count;
586   char *p;
587 
588   PetscFunctionBegin;
589   for (i=0,p=buf,left=len; i<n; i++) {
590     ierr = PetscSNPrintfCount(p,left,fmt,&count,x[i]);CHKERRQ(ierr);
591     if (count >= left) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_OUTOFRANGE,"Insufficient space in buffer");
592     left -= count;
593     p += count;
594     *p++ = ' ';
595   }
596   p[i ? 0 : -1] = 0;
597   PetscFunctionReturn(0);
598 }
599 
600 #undef __FUNCT__
601 #define __FUNCT__ "TSView_ARKIMEX"
602 static PetscErrorCode TSView_ARKIMEX(TS ts,PetscViewer viewer)
603 {
604   TS_ARKIMEX     *ark = (TS_ARKIMEX*)ts->data;
605   ARKTableau     tab = ark->tableau;
606   PetscBool      iascii;
607   PetscErrorCode ierr;
608 
609   PetscFunctionBegin;
610   ierr = PetscTypeCompare((PetscObject)viewer,PETSCVIEWERASCII,&iascii);CHKERRQ(ierr);
611   if (iascii) {
612     const TSARKIMEXType arktype;
613     char buf[512];
614     ierr = TSARKIMEXGetType(ts,&arktype);CHKERRQ(ierr);
615     ierr = PetscViewerASCIIPrintf(viewer,"  ARK IMEX %s\n",arktype);CHKERRQ(ierr);
616     ierr = PetscFormatRealArray(buf,sizeof buf,"% 8.6f",tab->s,tab->ct);CHKERRQ(ierr);
617     ierr = PetscViewerASCIIPrintf(viewer,"  Stiff abscissa       ct = %s\n",buf);CHKERRQ(ierr);
618     ierr = PetscFormatRealArray(buf,sizeof buf,"% 8.6f",tab->s,tab->c);CHKERRQ(ierr);
619     ierr = PetscViewerASCIIPrintf(viewer,"  Nonstiff abscissa     c = %s\n",buf);CHKERRQ(ierr);
620   }
621   ierr = SNESView(ts->snes,viewer);CHKERRQ(ierr);
622   PetscFunctionReturn(0);
623 }
624 
625 #undef __FUNCT__
626 #define __FUNCT__ "TSARKIMEXSetType"
627 /*@C
628   TSARKIMEXSetType - Set the type of ARK IMEX scheme
629 
630   Logically collective
631 
632   Input Parameter:
633 +  ts - timestepping context
634 -  arktype - type of ARK-IMEX scheme
635 
636   Level: intermediate
637 
638 .seealso: TSARKIMEXGetType()
639 @*/
640 PetscErrorCode TSARKIMEXSetType(TS ts,const TSARKIMEXType arktype)
641 {
642   PetscErrorCode ierr;
643 
644   PetscFunctionBegin;
645   PetscValidHeaderSpecific(ts,TS_CLASSID,1);
646   ierr = PetscTryMethod(ts,"TSARKIMEXSetType_C",(TS,const TSARKIMEXType),(ts,arktype));CHKERRQ(ierr);
647   PetscFunctionReturn(0);
648 }
649 
650 #undef __FUNCT__
651 #define __FUNCT__ "TSARKIMEXGetType"
652 /*@C
653   TSARKIMEXGetType - Get the type of ARK IMEX scheme
654 
655   Logically collective
656 
657   Input Parameter:
658 .  ts - timestepping context
659 
660   Output Parameter:
661 .  arktype - type of ARK-IMEX scheme
662 
663   Level: intermediate
664 
665 .seealso: TSARKIMEXGetType()
666 @*/
667 PetscErrorCode TSARKIMEXGetType(TS ts,const TSARKIMEXType *arktype)
668 {
669   PetscErrorCode ierr;
670 
671   PetscFunctionBegin;
672   PetscValidHeaderSpecific(ts,TS_CLASSID,1);
673   ierr = PetscUseMethod(ts,"TSARKIMEXGetType_C",(TS,const TSARKIMEXType*),(ts,arktype));CHKERRQ(ierr);
674   PetscFunctionReturn(0);
675 }
676 
677 #undef __FUNCT__
678 #define __FUNCT__ "TSARKIMEXSetFullyImplicit"
679 /*@C
680   TSARKIMEXSetFullyImplicit - Solve both parts of the equation implicitly
681 
682   Logically collective
683 
684   Input Parameter:
685 +  ts - timestepping context
686 -  flg - PETSC_TRUE for fully implicit
687 
688   Level: intermediate
689 
690 .seealso: TSARKIMEXGetType()
691 @*/
692 PetscErrorCode TSARKIMEXSetFullyImplicit(TS ts,PetscBool flg)
693 {
694   PetscErrorCode ierr;
695 
696   PetscFunctionBegin;
697   PetscValidHeaderSpecific(ts,TS_CLASSID,1);
698   ierr = PetscTryMethod(ts,"TSARKIMEXSetFullyImplicit_C",(TS,PetscBool),(ts,flg));CHKERRQ(ierr);
699   PetscFunctionReturn(0);
700 }
701 
702 EXTERN_C_BEGIN
703 #undef __FUNCT__
704 #define __FUNCT__ "TSARKIMEXGetType_ARKIMEX"
705 PetscErrorCode  TSARKIMEXGetType_ARKIMEX(TS ts,const TSARKIMEXType *arktype)
706 {
707   TS_ARKIMEX *ark = (TS_ARKIMEX*)ts->data;
708   PetscErrorCode ierr;
709 
710   PetscFunctionBegin;
711   if (!ark->tableau) {ierr = TSARKIMEXSetType(ts,TSARKIMEXDefault);CHKERRQ(ierr);}
712   *arktype = ark->tableau->name;
713   PetscFunctionReturn(0);
714 }
715 #undef __FUNCT__
716 #define __FUNCT__ "TSARKIMEXSetType_ARKIMEX"
717 PetscErrorCode  TSARKIMEXSetType_ARKIMEX(TS ts,const TSARKIMEXType arktype)
718 {
719   TS_ARKIMEX *ark = (TS_ARKIMEX*)ts->data;
720   PetscErrorCode ierr;
721   PetscBool match;
722   ARKTableauLink link;
723 
724   PetscFunctionBegin;
725   if (ark->tableau) {
726     ierr = PetscStrcmp(ark->tableau->name,arktype,&match);CHKERRQ(ierr);
727     if (match) PetscFunctionReturn(0);
728   }
729   for (link = ARKTableauList; link; link=link->next) {
730     ierr = PetscStrcmp(link->tab.name,arktype,&match);CHKERRQ(ierr);
731     if (match) {
732       ierr = TSReset_ARKIMEX(ts);CHKERRQ(ierr);
733       ark->tableau = &link->tab;
734       PetscFunctionReturn(0);
735     }
736   }
737   SETERRQ1(((PetscObject)ts)->comm,PETSC_ERR_ARG_UNKNOWN_TYPE,"Could not find '%s'",arktype);
738   PetscFunctionReturn(0);
739 }
740 #undef __FUNCT__
741 #define __FUNCT__ "TSARKIMEXSetFullyImplicit_ARKIMEX"
742 PetscErrorCode  TSARKIMEXSetFullyImplicit_ARKIMEX(TS ts,PetscBool flg)
743 {
744   TS_ARKIMEX *ark = (TS_ARKIMEX*)ts->data;
745 
746   PetscFunctionBegin;
747   ark->imex = (PetscBool)!flg;
748   PetscFunctionReturn(0);
749 }
750 EXTERN_C_END
751 
752 /* ------------------------------------------------------------ */
753 /*MC
754       TSARKIMEX - ODE solver using Additive Runge-Kutta IMEX schemes
755 
756   These methods are intended for problems with well-separated time scales, especially when a slow scale is strongly
757   nonlinear such that it is expensive to solve with a fully implicit method. The user should provide the stiff part
758   of the equation using TSSetIFunction() and the non-stiff part with TSSetRHSFunction().
759 
760   Notes:
761   The default is TSARKIMEX2E, it can be changed with TSARKIMEXSetType() or -ts_arkimex_type
762 
763   This method currently only works with ODE, for which the stiff part G(t,X,Xdot) has the form Xdot + Ghat(t,X).
764 
765   Level: beginner
766 
767 .seealso:  TSCreate(), TS, TSSetType(), TSARKIMEXSetType(), TSARKIMEXGetType(), TSARKIMEXSetFullyImplicit(), TSARKIMEX2D, TTSARKIMEX2E, TSARKIMEX3,
768            TSARKIMEX4, TSARKIMEX5, TSARKIMEXType, TSARKIMEXRegister()
769 
770 M*/
771 EXTERN_C_BEGIN
772 #undef __FUNCT__
773 #define __FUNCT__ "TSCreate_ARKIMEX"
774 PetscErrorCode  TSCreate_ARKIMEX(TS ts)
775 {
776   TS_ARKIMEX       *th;
777   PetscErrorCode ierr;
778 
779   PetscFunctionBegin;
780 #if !defined(PETSC_USE_DYNAMIC_LIBRARIES)
781   ierr = TSARKIMEXInitializePackage(PETSC_NULL);CHKERRQ(ierr);
782 #endif
783 
784   ts->ops->reset          = TSReset_ARKIMEX;
785   ts->ops->destroy        = TSDestroy_ARKIMEX;
786   ts->ops->view           = TSView_ARKIMEX;
787   ts->ops->setup          = TSSetUp_ARKIMEX;
788   ts->ops->step           = TSStep_ARKIMEX;
789   ts->ops->interpolate    = TSInterpolate_ARKIMEX;
790   ts->ops->setfromoptions = TSSetFromOptions_ARKIMEX;
791   ts->ops->snesfunction   = SNESTSFormFunction_ARKIMEX;
792   ts->ops->snesjacobian   = SNESTSFormJacobian_ARKIMEX;
793 
794   ierr = PetscNewLog(ts,TS_ARKIMEX,&th);CHKERRQ(ierr);
795   ts->data = (void*)th;
796   th->imex = PETSC_TRUE;
797 
798   ierr = PetscObjectComposeFunctionDynamic((PetscObject)ts,"TSARKIMEXGetType_C","TSARKIMEXGetType_ARKIMEX",TSARKIMEXGetType_ARKIMEX);CHKERRQ(ierr);
799   ierr = PetscObjectComposeFunctionDynamic((PetscObject)ts,"TSARKIMEXSetType_C","TSARKIMEXSetType_ARKIMEX",TSARKIMEXSetType_ARKIMEX);CHKERRQ(ierr);
800   ierr = PetscObjectComposeFunctionDynamic((PetscObject)ts,"TSARKIMEXSetFullyImplicit_C","TSARKIMEXSetFullyImplicit_ARKIMEX",TSARKIMEXSetFullyImplicit_ARKIMEX);CHKERRQ(ierr);
801   PetscFunctionReturn(0);
802 }
803 EXTERN_C_END
804