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