xref: /petsc/src/ts/impls/arkimex/arkimex.c (revision caa2e8a421fec08b263717f3d735599a28344997)
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   SNESConvergedReason snesreason;
374   PetscInt            i,j,its,lits;
375   PetscReal           next_time_step;
376   PetscReal           h,t;
377   PetscErrorCode      ierr;
378 
379   PetscFunctionBegin;
380   ierr = TSGetSNES(ts,&snes);CHKERRQ(ierr);
381   next_time_step = ts->time_step;
382   h = ts->time_step;
383   t = ts->ptime;
384 
385   for (i=0; i<s; i++) {
386     if (At[i*s+i] == 0) {           /* This stage is explicit */
387       ierr = VecCopy(ts->vec_sol,Y[i]);CHKERRQ(ierr);
388       for (j=0; j<i; j++) w[j] = -h*At[i*s+j];
389       ierr = VecMAXPY(Y[i],i,w,YdotI);CHKERRQ(ierr);
390       for (j=0; j<i; j++) w[j] = h*A[i*s+j];
391       ierr = VecMAXPY(Y[i],i,w,YdotRHS);CHKERRQ(ierr);
392     } else {
393       ark->stage_time = t + h*ct[i];
394       ark->shift = 1./(h*At[i*s+i]);
395       /* Affine part */
396       ierr = VecZeroEntries(W);CHKERRQ(ierr);
397       for (j=0; j<i; j++) w[j] = h*A[i*s+j];
398       ierr = VecMAXPY(W,i,w,YdotRHS);CHKERRQ(ierr);
399       /* Ydot = shift*(Y-Z) */
400       ierr = VecCopy(ts->vec_sol,Z);CHKERRQ(ierr);
401       for (j=0; j<i; j++) w[j] = -h*At[i*s+j];
402       ierr = VecMAXPY(Z,i,w,YdotI);CHKERRQ(ierr);
403       /* Initial guess taken from last stage */
404       ierr = VecCopy(i>0?Y[i-1]:ts->vec_sol,Y[i]);CHKERRQ(ierr);
405       ierr = SNESSolve(snes,W,Y[i]);CHKERRQ(ierr);
406       ierr = SNESGetIterationNumber(snes,&its);CHKERRQ(ierr);
407       ierr = SNESGetLinearSolveIterations(snes,&lits);CHKERRQ(ierr);
408       ierr = SNESGetConvergedReason(ts->snes,&snesreason);CHKERRQ(ierr);
409       ts->nonlinear_its += its; ts->linear_its += lits;
410       if (snesreason < 0 && ts->max_snes_failures > 0 && ++ts->num_snes_failures >= ts->max_snes_failures) {
411         ts->reason = TS_DIVERGED_NONLINEAR_SOLVE;
412         ierr = PetscInfo2(ts,"Step=%D, nonlinear solve solve failures %D greater than current TS allowed, stopping solve\n",ts->steps,ts->num_snes_failures);CHKERRQ(ierr);
413         PetscFunctionReturn(0);
414       }
415     }
416     ierr = VecZeroEntries(Ydot);CHKERRQ(ierr);
417     ierr = TSComputeIFunction(ts,t+h*ct[i],Y[i],Ydot,YdotI[i],ark->imex);CHKERRQ(ierr);
418     if (ark->imex) {
419       ierr = TSComputeRHSFunction(ts,t+h*c[i],Y[i],YdotRHS[i]);CHKERRQ(ierr);
420     } else {
421       ierr = VecZeroEntries(YdotRHS[i]);CHKERRQ(ierr);
422     }
423   }
424   for (j=0; j<s; j++) w[j] = -h*bt[j];
425   ierr = VecMAXPY(ts->vec_sol,s,w,YdotI);CHKERRQ(ierr);
426   for (j=0; j<s; j++) w[j] = h*b[j];
427   ierr = VecMAXPY(ts->vec_sol,s,w,YdotRHS);CHKERRQ(ierr);
428 
429   ts->ptime += ts->time_step;
430   ts->time_step = next_time_step;
431   ts->steps++;
432   PetscFunctionReturn(0);
433 }
434 
435 #undef __FUNCT__
436 #define __FUNCT__ "TSInterpolate_ARKIMEX"
437 static PetscErrorCode TSInterpolate_ARKIMEX(TS ts,PetscReal itime,Vec X)
438 {
439   TS_ARKIMEX *ark = (TS_ARKIMEX*)ts->data;
440   PetscInt s = ark->tableau->s,pinterp = ark->tableau->pinterp,i,j;
441   PetscReal tt,t = (itime - ts->ptime)/ts->time_step + 1; /* In the interval [0,1] */
442   PetscScalar *bt,*b;
443   const PetscReal *Bt = ark->tableau->binterpt,*B = ark->tableau->binterp;
444   PetscErrorCode ierr;
445 
446   PetscFunctionBegin;
447   if (!Bt || !B) SETERRQ1(((PetscObject)ts)->comm,PETSC_ERR_SUP,"TSARKIMEX %s does not have an interpolation formula",ark->tableau->name);
448   ierr = PetscMalloc2(s,PetscScalar,&bt,s,PetscScalar,&b);CHKERRQ(ierr);
449   for (i=0; i<s; i++) bt[i] = b[i] = 0;
450   for (j=0,tt=t; j<pinterp; j++,tt*=t) {
451     for (i=0; i<s; i++) {
452       bt[i] += ts->time_step * Bt[i*pinterp+j] * tt * -1.0;
453       b[i]  += ts->time_step * B[i*pinterp+j] * tt;
454     }
455   }
456   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");
457   ierr = VecCopy(ark->Y[0],X);CHKERRQ(ierr);
458   ierr = VecMAXPY(X,s,bt,ark->YdotI);CHKERRQ(ierr);
459   ierr = VecMAXPY(X,s,b,ark->YdotRHS);CHKERRQ(ierr);
460   ierr = PetscFree2(bt,b);CHKERRQ(ierr);
461   PetscFunctionReturn(0);
462 }
463 
464 /*------------------------------------------------------------*/
465 #undef __FUNCT__
466 #define __FUNCT__ "TSReset_ARKIMEX"
467 static PetscErrorCode TSReset_ARKIMEX(TS ts)
468 {
469   TS_ARKIMEX      *ark = (TS_ARKIMEX*)ts->data;
470   PetscInt        s;
471   PetscErrorCode  ierr;
472 
473   PetscFunctionBegin;
474   if (!ark->tableau) PetscFunctionReturn(0);
475    s = ark->tableau->s;
476   ierr = VecDestroyVecs(s,&ark->Y);CHKERRQ(ierr);
477   ierr = VecDestroyVecs(s,&ark->YdotI);CHKERRQ(ierr);
478   ierr = VecDestroyVecs(s,&ark->YdotRHS);CHKERRQ(ierr);
479   ierr = VecDestroy(&ark->Ydot);CHKERRQ(ierr);
480   ierr = VecDestroy(&ark->Work);CHKERRQ(ierr);
481   ierr = VecDestroy(&ark->Z);CHKERRQ(ierr);
482   ierr = PetscFree(ark->work);CHKERRQ(ierr);
483   PetscFunctionReturn(0);
484 }
485 
486 #undef __FUNCT__
487 #define __FUNCT__ "TSDestroy_ARKIMEX"
488 static PetscErrorCode TSDestroy_ARKIMEX(TS ts)
489 {
490   PetscErrorCode  ierr;
491 
492   PetscFunctionBegin;
493   ierr = TSReset_ARKIMEX(ts);CHKERRQ(ierr);
494   ierr = PetscFree(ts->data);CHKERRQ(ierr);
495   ierr = PetscObjectComposeFunctionDynamic((PetscObject)ts,"TSARKIMEXGetType_C","",PETSC_NULL);CHKERRQ(ierr);
496   ierr = PetscObjectComposeFunctionDynamic((PetscObject)ts,"TSARKIMEXSetType_C","",PETSC_NULL);CHKERRQ(ierr);
497   PetscFunctionReturn(0);
498 }
499 
500 /*
501   This defines the nonlinear equation that is to be solved with SNES
502   G(U) = F[t0+Theta*dt, U, (U-U0)*shift] = 0
503 */
504 #undef __FUNCT__
505 #define __FUNCT__ "SNESTSFormFunction_ARKIMEX"
506 static PetscErrorCode SNESTSFormFunction_ARKIMEX(SNES snes,Vec X,Vec F,TS ts)
507 {
508   TS_ARKIMEX     *ark = (TS_ARKIMEX*)ts->data;
509   PetscErrorCode ierr;
510 
511   PetscFunctionBegin;
512   ierr = VecAXPBYPCZ(ark->Ydot,-ark->shift,ark->shift,0,ark->Z,X);CHKERRQ(ierr); /* Ydot = shift*(X-Z) */
513   ierr = TSComputeIFunction(ts,ark->stage_time,X,ark->Ydot,F,ark->imex);CHKERRQ(ierr);
514   PetscFunctionReturn(0);
515 }
516 
517 #undef __FUNCT__
518 #define __FUNCT__ "SNESTSFormJacobian_ARKIMEX"
519 static PetscErrorCode SNESTSFormJacobian_ARKIMEX(SNES snes,Vec X,Mat *A,Mat *B,MatStructure *str,TS ts)
520 {
521   TS_ARKIMEX       *ark = (TS_ARKIMEX*)ts->data;
522   PetscErrorCode ierr;
523 
524   PetscFunctionBegin;
525   /* ark->Ydot has already been computed in SNESTSFormFunction_ARKIMEX (SNES guarantees this) */
526   ierr = TSComputeIJacobian(ts,ark->stage_time,X,ark->Ydot,ark->shift,A,B,str,PETSC_TRUE);CHKERRQ(ierr);
527   PetscFunctionReturn(0);
528 }
529 
530 #undef __FUNCT__
531 #define __FUNCT__ "TSSetUp_ARKIMEX"
532 static PetscErrorCode TSSetUp_ARKIMEX(TS ts)
533 {
534   TS_ARKIMEX     *ark = (TS_ARKIMEX*)ts->data;
535   ARKTableau     tab  = ark->tableau;
536   PetscInt       s = tab->s;
537   PetscErrorCode ierr;
538 
539   PetscFunctionBegin;
540   if (!ark->tableau) {
541     ierr = TSARKIMEXSetType(ts,TSARKIMEXDefault);CHKERRQ(ierr);
542   }
543   ierr = VecDuplicateVecs(ts->vec_sol,s,&ark->Y);CHKERRQ(ierr);
544   ierr = VecDuplicateVecs(ts->vec_sol,s,&ark->YdotI);CHKERRQ(ierr);
545   ierr = VecDuplicateVecs(ts->vec_sol,s,&ark->YdotRHS);CHKERRQ(ierr);
546   ierr = VecDuplicate(ts->vec_sol,&ark->Ydot);CHKERRQ(ierr);
547   ierr = VecDuplicate(ts->vec_sol,&ark->Work);CHKERRQ(ierr);
548   ierr = VecDuplicate(ts->vec_sol,&ark->Z);CHKERRQ(ierr);
549   ierr = PetscMalloc(s*sizeof(ark->work[0]),&ark->work);CHKERRQ(ierr);
550   PetscFunctionReturn(0);
551 }
552 /*------------------------------------------------------------*/
553 
554 #undef __FUNCT__
555 #define __FUNCT__ "TSSetFromOptions_ARKIMEX"
556 static PetscErrorCode TSSetFromOptions_ARKIMEX(TS ts)
557 {
558   TS_ARKIMEX     *ark = (TS_ARKIMEX*)ts->data;
559   PetscErrorCode ierr;
560   char           arktype[256];
561 
562   PetscFunctionBegin;
563   ierr = PetscOptionsHead("ARKIMEX ODE solver options");CHKERRQ(ierr);
564   {
565     ARKTableauLink link;
566     PetscInt count,choice;
567     PetscBool flg;
568     const char **namelist;
569     ierr = PetscStrncpy(arktype,TSARKIMEXDefault,sizeof arktype);CHKERRQ(ierr);
570     for (link=ARKTableauList,count=0; link; link=link->next,count++) ;
571     ierr = PetscMalloc(count*sizeof(char*),&namelist);CHKERRQ(ierr);
572     for (link=ARKTableauList,count=0; link; link=link->next,count++) namelist[count] = link->tab.name;
573     ierr = PetscOptionsEList("-ts_arkimex_type","Family of ARK IMEX method","TSARKIMEXSetType",(const char*const*)namelist,count,arktype,&choice,&flg);CHKERRQ(ierr);
574     ierr = TSARKIMEXSetType(ts,flg ? namelist[choice] : arktype);CHKERRQ(ierr);
575     ierr = PetscFree(namelist);CHKERRQ(ierr);
576     flg = (PetscBool)!ark->imex;
577     ierr = PetscOptionsBool("-ts_arkimex_fully_implicit","Solve the problem fully implicitly","TSARKIMEXSetFullyImplicit",flg,&flg,PETSC_NULL);CHKERRQ(ierr);
578     ark->imex = (PetscBool)!flg;
579     ierr = SNESSetFromOptions(ts->snes);CHKERRQ(ierr);
580   }
581   ierr = PetscOptionsTail();CHKERRQ(ierr);
582   PetscFunctionReturn(0);
583 }
584 
585 #undef __FUNCT__
586 #define __FUNCT__ "PetscFormatRealArray"
587 static PetscErrorCode PetscFormatRealArray(char buf[],size_t len,const char *fmt,PetscInt n,const PetscReal x[])
588 {
589   PetscErrorCode ierr;
590   int i,left,count;
591   char *p;
592 
593   PetscFunctionBegin;
594   for (i=0,p=buf,left=(int)len; i<n; i++) {
595     ierr = PetscSNPrintf(p,left,fmt,&count,x[i]);CHKERRQ(ierr);
596     if (count >= left) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_OUTOFRANGE,"Insufficient space in buffer");
597     left -= count;
598     p += count;
599     *p++ = ' ';
600   }
601   p[i ? 0 : -1] = 0;
602   PetscFunctionReturn(0);
603 }
604 
605 #undef __FUNCT__
606 #define __FUNCT__ "TSView_ARKIMEX"
607 static PetscErrorCode TSView_ARKIMEX(TS ts,PetscViewer viewer)
608 {
609   TS_ARKIMEX     *ark = (TS_ARKIMEX*)ts->data;
610   ARKTableau     tab = ark->tableau;
611   PetscBool      iascii;
612   PetscErrorCode ierr;
613 
614   PetscFunctionBegin;
615   ierr = PetscTypeCompare((PetscObject)viewer,PETSCVIEWERASCII,&iascii);CHKERRQ(ierr);
616   if (iascii) {
617     const TSARKIMEXType arktype;
618     char buf[512];
619     ierr = TSARKIMEXGetType(ts,&arktype);CHKERRQ(ierr);
620     ierr = PetscViewerASCIIPrintf(viewer,"  ARK IMEX %s\n",arktype);CHKERRQ(ierr);
621     ierr = PetscFormatRealArray(buf,sizeof buf,"% 8.6f",tab->s,tab->ct);CHKERRQ(ierr);
622     ierr = PetscViewerASCIIPrintf(viewer,"  Stiff abscissa       ct = %s\n",buf);CHKERRQ(ierr);
623     ierr = PetscFormatRealArray(buf,sizeof buf,"% 8.6f",tab->s,tab->c);CHKERRQ(ierr);
624     ierr = PetscViewerASCIIPrintf(viewer,"  Nonstiff abscissa     c = %s\n",buf);CHKERRQ(ierr);
625   }
626   ierr = SNESView(ts->snes,viewer);CHKERRQ(ierr);
627   PetscFunctionReturn(0);
628 }
629 
630 #undef __FUNCT__
631 #define __FUNCT__ "TSARKIMEXSetType"
632 /*@C
633   TSARKIMEXSetType - Set the type of ARK IMEX scheme
634 
635   Logically collective
636 
637   Input Parameter:
638 +  ts - timestepping context
639 -  arktype - type of ARK-IMEX scheme
640 
641   Level: intermediate
642 
643 .seealso: TSARKIMEXGetType()
644 @*/
645 PetscErrorCode TSARKIMEXSetType(TS ts,const TSARKIMEXType arktype)
646 {
647   PetscErrorCode ierr;
648 
649   PetscFunctionBegin;
650   PetscValidHeaderSpecific(ts,TS_CLASSID,1);
651   ierr = PetscTryMethod(ts,"TSARKIMEXSetType_C",(TS,const TSARKIMEXType),(ts,arktype));CHKERRQ(ierr);
652   PetscFunctionReturn(0);
653 }
654 
655 #undef __FUNCT__
656 #define __FUNCT__ "TSARKIMEXGetType"
657 /*@C
658   TSARKIMEXGetType - Get the type of ARK IMEX scheme
659 
660   Logically collective
661 
662   Input Parameter:
663 .  ts - timestepping context
664 
665   Output Parameter:
666 .  arktype - type of ARK-IMEX scheme
667 
668   Level: intermediate
669 
670 .seealso: TSARKIMEXGetType()
671 @*/
672 PetscErrorCode TSARKIMEXGetType(TS ts,const TSARKIMEXType *arktype)
673 {
674   PetscErrorCode ierr;
675 
676   PetscFunctionBegin;
677   PetscValidHeaderSpecific(ts,TS_CLASSID,1);
678   ierr = PetscUseMethod(ts,"TSARKIMEXGetType_C",(TS,const TSARKIMEXType*),(ts,arktype));CHKERRQ(ierr);
679   PetscFunctionReturn(0);
680 }
681 
682 #undef __FUNCT__
683 #define __FUNCT__ "TSARKIMEXSetFullyImplicit"
684 /*@C
685   TSARKIMEXSetFullyImplicit - Solve both parts of the equation implicitly
686 
687   Logically collective
688 
689   Input Parameter:
690 +  ts - timestepping context
691 -  flg - PETSC_TRUE for fully implicit
692 
693   Level: intermediate
694 
695 .seealso: TSARKIMEXGetType()
696 @*/
697 PetscErrorCode TSARKIMEXSetFullyImplicit(TS ts,PetscBool flg)
698 {
699   PetscErrorCode ierr;
700 
701   PetscFunctionBegin;
702   PetscValidHeaderSpecific(ts,TS_CLASSID,1);
703   ierr = PetscTryMethod(ts,"TSARKIMEXSetFullyImplicit_C",(TS,PetscBool),(ts,flg));CHKERRQ(ierr);
704   PetscFunctionReturn(0);
705 }
706 
707 EXTERN_C_BEGIN
708 #undef __FUNCT__
709 #define __FUNCT__ "TSARKIMEXGetType_ARKIMEX"
710 PetscErrorCode  TSARKIMEXGetType_ARKIMEX(TS ts,const TSARKIMEXType *arktype)
711 {
712   TS_ARKIMEX *ark = (TS_ARKIMEX*)ts->data;
713   PetscErrorCode ierr;
714 
715   PetscFunctionBegin;
716   if (!ark->tableau) {ierr = TSARKIMEXSetType(ts,TSARKIMEXDefault);CHKERRQ(ierr);}
717   *arktype = ark->tableau->name;
718   PetscFunctionReturn(0);
719 }
720 #undef __FUNCT__
721 #define __FUNCT__ "TSARKIMEXSetType_ARKIMEX"
722 PetscErrorCode  TSARKIMEXSetType_ARKIMEX(TS ts,const TSARKIMEXType arktype)
723 {
724   TS_ARKIMEX *ark = (TS_ARKIMEX*)ts->data;
725   PetscErrorCode ierr;
726   PetscBool match;
727   ARKTableauLink link;
728 
729   PetscFunctionBegin;
730   if (ark->tableau) {
731     ierr = PetscStrcmp(ark->tableau->name,arktype,&match);CHKERRQ(ierr);
732     if (match) PetscFunctionReturn(0);
733   }
734   for (link = ARKTableauList; link; link=link->next) {
735     ierr = PetscStrcmp(link->tab.name,arktype,&match);CHKERRQ(ierr);
736     if (match) {
737       ierr = TSReset_ARKIMEX(ts);CHKERRQ(ierr);
738       ark->tableau = &link->tab;
739       PetscFunctionReturn(0);
740     }
741   }
742   SETERRQ1(((PetscObject)ts)->comm,PETSC_ERR_ARG_UNKNOWN_TYPE,"Could not find '%s'",arktype);
743   PetscFunctionReturn(0);
744 }
745 #undef __FUNCT__
746 #define __FUNCT__ "TSARKIMEXSetFullyImplicit_ARKIMEX"
747 PetscErrorCode  TSARKIMEXSetFullyImplicit_ARKIMEX(TS ts,PetscBool flg)
748 {
749   TS_ARKIMEX *ark = (TS_ARKIMEX*)ts->data;
750 
751   PetscFunctionBegin;
752   ark->imex = (PetscBool)!flg;
753   PetscFunctionReturn(0);
754 }
755 EXTERN_C_END
756 
757 /* ------------------------------------------------------------ */
758 /*MC
759       TSARKIMEX - ODE solver using Additive Runge-Kutta IMEX schemes
760 
761   These methods are intended for problems with well-separated time scales, especially when a slow scale is strongly
762   nonlinear such that it is expensive to solve with a fully implicit method. The user should provide the stiff part
763   of the equation using TSSetIFunction() and the non-stiff part with TSSetRHSFunction().
764 
765   Notes:
766   The default is TSARKIMEX2E, it can be changed with TSARKIMEXSetType() or -ts_arkimex_type
767 
768   This method currently only works with ODE, for which the stiff part G(t,X,Xdot) has the form Xdot + Ghat(t,X).
769 
770   Level: beginner
771 
772 .seealso:  TSCreate(), TS, TSSetType(), TSARKIMEXSetType(), TSARKIMEXGetType(), TSARKIMEXSetFullyImplicit(), TSARKIMEX2D, TTSARKIMEX2E, TSARKIMEX3,
773            TSARKIMEX4, TSARKIMEX5, TSARKIMEXType, TSARKIMEXRegister()
774 
775 M*/
776 EXTERN_C_BEGIN
777 #undef __FUNCT__
778 #define __FUNCT__ "TSCreate_ARKIMEX"
779 PetscErrorCode  TSCreate_ARKIMEX(TS ts)
780 {
781   TS_ARKIMEX       *th;
782   PetscErrorCode ierr;
783 
784   PetscFunctionBegin;
785 #if !defined(PETSC_USE_DYNAMIC_LIBRARIES)
786   ierr = TSARKIMEXInitializePackage(PETSC_NULL);CHKERRQ(ierr);
787 #endif
788 
789   ts->ops->reset          = TSReset_ARKIMEX;
790   ts->ops->destroy        = TSDestroy_ARKIMEX;
791   ts->ops->view           = TSView_ARKIMEX;
792   ts->ops->setup          = TSSetUp_ARKIMEX;
793   ts->ops->step           = TSStep_ARKIMEX;
794   ts->ops->interpolate    = TSInterpolate_ARKIMEX;
795   ts->ops->setfromoptions = TSSetFromOptions_ARKIMEX;
796   ts->ops->snesfunction   = SNESTSFormFunction_ARKIMEX;
797   ts->ops->snesjacobian   = SNESTSFormJacobian_ARKIMEX;
798 
799   ierr = PetscNewLog(ts,TS_ARKIMEX,&th);CHKERRQ(ierr);
800   ts->data = (void*)th;
801   th->imex = PETSC_TRUE;
802 
803   ierr = PetscObjectComposeFunctionDynamic((PetscObject)ts,"TSARKIMEXGetType_C","TSARKIMEXGetType_ARKIMEX",TSARKIMEXGetType_ARKIMEX);CHKERRQ(ierr);
804   ierr = PetscObjectComposeFunctionDynamic((PetscObject)ts,"TSARKIMEXSetType_C","TSARKIMEXSetType_ARKIMEX",TSARKIMEXSetType_ARKIMEX);CHKERRQ(ierr);
805   ierr = PetscObjectComposeFunctionDynamic((PetscObject)ts,"TSARKIMEXSetFullyImplicit_C","TSARKIMEXSetFullyImplicit_ARKIMEX",TSARKIMEXSetFullyImplicit_ARKIMEX);CHKERRQ(ierr);
806   PetscFunctionReturn(0);
807 }
808 EXTERN_C_END
809