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