Actual source code: irk.c

  1: /*
  2:   Code for timestepping with implicit Runge-Kutta method

  4:   Notes:
  5:   The general system is written as

  7:   F(t,U,Udot) = 0

  9: */
 10: #include <petsc/private/tsimpl.h>
 11: #include <petscdm.h>
 12: #include <petscdt.h>

 14: static TSIRKType         TSIRKDefault = TSIRKGAUSS;
 15: static PetscBool         TSIRKRegisterAllCalled;
 16: static PetscBool         TSIRKPackageInitialized;
 17: static PetscFunctionList TSIRKList;

 19: struct _IRKTableau{
 20:   PetscReal   *A,*b,*c;
 21:   PetscScalar *A_inv,*A_inv_rowsum,*I_s;
 22:   PetscReal   *binterp;   /* Dense output formula */
 23: };

 25: typedef struct _IRKTableau *IRKTableau;

 27: typedef struct {
 28:   char         *method_name;
 29:   PetscInt     order;            /* Classical approximation order of the method */
 30:   PetscInt     nstages;          /* Number of stages */
 31:   PetscBool    stiffly_accurate;
 32:   PetscInt     pinterp;          /* Interpolation order */
 33:   IRKTableau   tableau;
 34:   Vec          U0;               /* Backup vector */
 35:   Vec          Z;                /* Combined stage vector */
 36:   Vec          *Y;               /* States computed during the step */
 37:   Vec          Ydot;             /* Work vector holding time derivatives during residual evaluation */
 38:   Vec          U;                /* U is used to compute Ydot = shift(Y-U) */
 39:   Vec          *YdotI;           /* Work vectors to hold the residual evaluation */
 40:   Mat          TJ;               /* KAIJ matrix for the Jacobian of the combined system */
 41:   PetscScalar  *work;            /* Scalar work */
 42:   TSStepStatus status;
 43:   PetscBool    rebuild_completion;
 44:   PetscReal    ccfl;
 45: } TS_IRK;

 47: /*@C
 48:    TSIRKTableauCreate - create the tableau for TSIRK and provide the entries

 50:    Not Collective

 52:    Input Parameters:
 53: +  ts - timestepping context
 54: .  nstages - number of stages, this is the dimension of the matrices below
 55: .  A - stage coefficients (dimension nstages*nstages, row-major)
 56: .  b - step completion table (dimension nstages)
 57: .  c - abscissa (dimension nstages)
 58: .  binterp - coefficients of the interpolation formula (dimension nstages)
 59: .  A_inv - inverse of A (dimension nstages*nstages, row-major)
 60: .  A_inv_rowsum - row sum of the inverse of A (dimension nstages)
 61: -  I_s - identity matrix (dimension nstages*nstages)

 63:    Level: advanced

 65: .seealso: TSIRK, TSIRKRegister()
 66: @*/
 67: PetscErrorCode TSIRKTableauCreate(TS ts,PetscInt nstages,const PetscReal *A,const PetscReal *b,const PetscReal *c,const PetscReal *binterp,const PetscScalar *A_inv,const PetscScalar *A_inv_rowsum,const PetscScalar *I_s)
 68: {
 69:   TS_IRK         *irk = (TS_IRK*)ts->data;
 70:   IRKTableau     tab = irk->tableau;

 72:   irk->order = nstages;
 73:   PetscMalloc3(PetscSqr(nstages),&tab->A,PetscSqr(nstages),&tab->A_inv,PetscSqr(nstages),&tab->I_s);
 74:   PetscMalloc4(nstages,&tab->b,nstages,&tab->c,nstages,&tab->binterp,nstages,&tab->A_inv_rowsum);
 75:   PetscArraycpy(tab->A,A,PetscSqr(nstages));
 76:   PetscArraycpy(tab->b,b,nstages);
 77:   PetscArraycpy(tab->c,c,nstages);
 78:   /* optional coefficient arrays */
 79:   if (binterp) {
 80:     PetscArraycpy(tab->binterp,binterp,nstages);
 81:   }
 82:   if (A_inv) {
 83:     PetscArraycpy(tab->A_inv,A_inv,PetscSqr(nstages));
 84:   }
 85:   if (A_inv_rowsum) {
 86:     PetscArraycpy(tab->A_inv_rowsum,A_inv_rowsum,nstages);
 87:   }
 88:   if (I_s) {
 89:     PetscArraycpy(tab->I_s,I_s,PetscSqr(nstages));
 90:   }
 91:   return 0;
 92: }

 94: /* Arrays should be freed with PetscFree3(A,b,c) */
 95: static PetscErrorCode TSIRKCreate_Gauss(TS ts)
 96: {
 97:   PetscInt       nstages;
 98:   PetscReal      *gauss_A_real,*gauss_b,*b,*gauss_c;
 99:   PetscScalar    *gauss_A,*gauss_A_inv,*gauss_A_inv_rowsum,*I_s;
100:   PetscScalar    *G0,*G1;
101:   PetscInt       i,j;
102:   Mat            G0mat,G1mat,Amat;

104:   TSIRKGetNumStages(ts,&nstages);
105:   PetscMalloc3(PetscSqr(nstages),&gauss_A_real,nstages,&gauss_b,nstages,&gauss_c);
106:   PetscMalloc4(PetscSqr(nstages),&gauss_A,PetscSqr(nstages),&gauss_A_inv,nstages,&gauss_A_inv_rowsum,PetscSqr(nstages),&I_s);
107:   PetscMalloc3(nstages,&b,PetscSqr(nstages),&G0,PetscSqr(nstages),&G1);
108:   PetscDTGaussQuadrature(nstages,0.,1.,gauss_c,b);
109:   for (i=0; i<nstages; i++) gauss_b[i] = b[i]; /* copy to possibly-complex array */

111:   /* A^T = G0^{-1} G1 */
112:   for (i=0; i<nstages; i++) {
113:     for (j=0; j<nstages; j++) {
114:       G0[i*nstages+j] = PetscPowRealInt(gauss_c[i],j);
115:       G1[i*nstages+j] = PetscPowRealInt(gauss_c[i],j+1)/(j+1);
116:     }
117:   }
118:   /* The arrays above are row-aligned, but we create dense matrices as the transpose */
119:   MatCreateSeqDense(PETSC_COMM_SELF,nstages,nstages,G0,&G0mat);
120:   MatCreateSeqDense(PETSC_COMM_SELF,nstages,nstages,G1,&G1mat);
121:   MatCreateSeqDense(PETSC_COMM_SELF,nstages,nstages,gauss_A,&Amat);
122:   MatLUFactor(G0mat,NULL,NULL,NULL);
123:   MatMatSolve(G0mat,G1mat,Amat);
124:   MatTranspose(Amat,MAT_INPLACE_MATRIX,&Amat);
125:   for (i=0; i<nstages; i++)
126:     for (j=0; j<nstages; j++)
127:       gauss_A_real[i*nstages+j] = PetscRealPart(gauss_A[i*nstages+j]);

129:   MatDestroy(&G0mat);
130:   MatDestroy(&G1mat);
131:   MatDestroy(&Amat);
132:   PetscFree3(b,G0,G1);

134:   {/* Invert A */
135:     /* PETSc does not provide a routine to calculate the inverse of a general matrix.
136:      * To get the inverse of A, we form a sequential BAIJ matrix from it, consisting of a single block with block size
137:      * equal to the dimension of A, and then use MatInvertBlockDiagonal(). */
138:     Mat               A_baij;
139:     PetscInt          idxm[1]={0},idxn[1]={0};
140:     const PetscScalar *A_inv;

142:     MatCreateSeqBAIJ(PETSC_COMM_SELF,nstages,nstages,nstages,1,NULL,&A_baij);
143:     MatSetOption(A_baij,MAT_ROW_ORIENTED,PETSC_FALSE);
144:     MatSetValuesBlocked(A_baij,1,idxm,1,idxn,gauss_A,INSERT_VALUES);
145:     MatAssemblyBegin(A_baij,MAT_FINAL_ASSEMBLY);
146:     MatAssemblyEnd(A_baij,MAT_FINAL_ASSEMBLY);
147:     MatInvertBlockDiagonal(A_baij,&A_inv);
148:     PetscMemcpy(gauss_A_inv,A_inv,nstages*nstages*sizeof(PetscScalar));
149:     MatDestroy(&A_baij);
150:   }

152:   /* Compute row sums A_inv_rowsum and identity I_s */
153:   for (i=0; i<nstages; i++) {
154:     gauss_A_inv_rowsum[i] = 0;
155:     for (j=0; j<nstages; j++) {
156:       gauss_A_inv_rowsum[i] += gauss_A_inv[i+nstages*j];
157:       I_s[i+nstages*j] = 1.*(i == j);
158:     }
159:   }
160:   TSIRKTableauCreate(ts,nstages,gauss_A_real,gauss_b,gauss_c,NULL,gauss_A_inv,gauss_A_inv_rowsum,I_s);
161:   PetscFree3(gauss_A_real,gauss_b,gauss_c);
162:   PetscFree4(gauss_A,gauss_A_inv,gauss_A_inv_rowsum,I_s);
163:   return 0;
164: }

166: /*@C
167:    TSIRKRegister -  adds a TSIRK implementation

169:    Not Collective

171:    Input Parameters:
172: +  sname - name of user-defined IRK scheme
173: -  function - function to create method context

175:    Notes:
176:    TSIRKRegister() may be called multiple times to add several user-defined families.

178:    Sample usage:
179: .vb
180:    TSIRKRegister("my_scheme",MySchemeCreate);
181: .ve

183:    Then, your scheme can be chosen with the procedural interface via
184: $     TSIRKSetType(ts,"my_scheme")
185:    or at runtime via the option
186: $     -ts_irk_type my_scheme

188:    Level: advanced

190: .seealso: TSIRKRegisterAll()
191: @*/
192: PetscErrorCode TSIRKRegister(const char sname[],PetscErrorCode (*function)(TS))
193: {
194:   TSIRKInitializePackage();
195:   PetscFunctionListAdd(&TSIRKList,sname,function);
196:   return 0;
197: }

199: /*@C
200:   TSIRKRegisterAll - Registers all of the implicit Runge-Kutta methods in TSIRK

202:   Not Collective, but should be called by all processes which will need the schemes to be registered

204:   Level: advanced

206: .seealso:  TSIRKRegisterDestroy()
207: @*/
208: PetscErrorCode TSIRKRegisterAll(void)
209: {
210:   if (TSIRKRegisterAllCalled) return 0;
211:   TSIRKRegisterAllCalled = PETSC_TRUE;

213:   TSIRKRegister(TSIRKGAUSS,TSIRKCreate_Gauss);
214:   return 0;
215: }

217: /*@C
218:    TSIRKRegisterDestroy - Frees the list of schemes that were registered by TSIRKRegister().

220:    Not Collective

222:    Level: advanced

224: .seealso: TSIRKRegister(), TSIRKRegisterAll()
225: @*/
226: PetscErrorCode TSIRKRegisterDestroy(void)
227: {
228:   TSIRKRegisterAllCalled = PETSC_FALSE;
229:   return 0;
230: }

232: /*@C
233:   TSIRKInitializePackage - This function initializes everything in the TSIRK package. It is called
234:   from TSInitializePackage().

236:   Level: developer

238: .seealso: PetscInitialize()
239: @*/
240: PetscErrorCode TSIRKInitializePackage(void)
241: {
242:   if (TSIRKPackageInitialized) return 0;
243:   TSIRKPackageInitialized = PETSC_TRUE;
244:   TSIRKRegisterAll();
245:   PetscRegisterFinalize(TSIRKFinalizePackage);
246:   return 0;
247: }

249: /*@C
250:   TSIRKFinalizePackage - This function destroys everything in the TSIRK package. It is
251:   called from PetscFinalize().

253:   Level: developer

255: .seealso: PetscFinalize()
256: @*/
257: PetscErrorCode TSIRKFinalizePackage(void)
258: {
259:   PetscFunctionListDestroy(&TSIRKList);
260:   TSIRKPackageInitialized = PETSC_FALSE;
261:   return 0;
262: }

264: /*
265:  This function can be called before or after ts->vec_sol has been updated.
266: */
267: static PetscErrorCode TSEvaluateStep_IRK(TS ts,PetscInt order,Vec U,PetscBool *done)
268: {
269:   TS_IRK         *irk = (TS_IRK*)ts->data;
270:   IRKTableau     tab = irk->tableau;
271:   Vec            *YdotI = irk->YdotI;
272:   PetscScalar    *w = irk->work;
273:   PetscReal      h;
274:   PetscInt       j;

276:   switch (irk->status) {
277:   case TS_STEP_INCOMPLETE:
278:   case TS_STEP_PENDING:
279:     h = ts->time_step; break;
280:   case TS_STEP_COMPLETE:
281:     h = ts->ptime - ts->ptime_prev; break;
282:   default: SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_PLIB,"Invalid TSStepStatus");
283:   }

285:   VecCopy(ts->vec_sol,U);
286:   for (j=0; j<irk->nstages; j++) w[j] = h*tab->b[j];
287:   VecMAXPY(U,irk->nstages,w,YdotI);
288:   return 0;
289: }

291: static PetscErrorCode TSRollBack_IRK(TS ts)
292: {
293:   TS_IRK         *irk = (TS_IRK*)ts->data;

295:   VecCopy(irk->U0,ts->vec_sol);
296:   return 0;
297: }

299: static PetscErrorCode TSStep_IRK(TS ts)
300: {
301:   TS_IRK          *irk = (TS_IRK*)ts->data;
302:   IRKTableau      tab = irk->tableau;
303:   PetscScalar     *A_inv = tab->A_inv,*A_inv_rowsum = tab->A_inv_rowsum;
304:   const PetscInt  nstages = irk->nstages;
305:   SNES            snes;
306:   PetscInt        i,j,its,lits,bs;
307:   TSAdapt         adapt;
308:   PetscInt        rejections = 0;
309:   PetscBool       accept = PETSC_TRUE;
310:   PetscReal       next_time_step = ts->time_step;

312:   if (!ts->steprollback) {
313:     VecCopy(ts->vec_sol,irk->U0);
314:   }
315:   VecGetBlockSize(ts->vec_sol,&bs);
316:   for (i=0; i<nstages; i++) {
317:     VecStrideScatter(ts->vec_sol,i*bs,irk->Z,INSERT_VALUES);
318:   }

320:   irk->status = TS_STEP_INCOMPLETE;
321:   while (!ts->reason && irk->status != TS_STEP_COMPLETE) {
322:     VecCopy(ts->vec_sol,irk->U);
323:     TSGetSNES(ts,&snes);
324:     SNESSolve(snes,NULL,irk->Z);
325:     SNESGetIterationNumber(snes,&its);
326:     SNESGetLinearSolveIterations(snes,&lits);
327:     ts->snes_its += its; ts->ksp_its += lits;
328:     VecStrideGatherAll(irk->Z,irk->Y,INSERT_VALUES);
329:     for (i=0; i<nstages; i++) {
330:       VecZeroEntries(irk->YdotI[i]);
331:       for (j=0; j<nstages; j++) {
332:         VecAXPY(irk->YdotI[i],A_inv[i+j*nstages]/ts->time_step,irk->Y[j]);
333:       }
334:       VecAXPY(irk->YdotI[i],-A_inv_rowsum[i]/ts->time_step,irk->U);
335:     }
336:     irk->status = TS_STEP_INCOMPLETE;
337:     TSEvaluateStep_IRK(ts,irk->order,ts->vec_sol,NULL);
338:     irk->status = TS_STEP_PENDING;
339:     TSGetAdapt(ts,&adapt);
340:     TSAdaptChoose(adapt,ts,ts->time_step,NULL,&next_time_step,&accept);
341:     irk->status = accept ? TS_STEP_COMPLETE : TS_STEP_INCOMPLETE;
342:     if (!accept) {
343:       TSRollBack_IRK(ts);
344:       ts->time_step = next_time_step;
345:       goto reject_step;
346:     }

348:     ts->ptime += ts->time_step;
349:     ts->time_step = next_time_step;
350:     break;
351:   reject_step:
352:     ts->reject++; accept = PETSC_FALSE;
353:     if (!ts->reason && ++rejections > ts->max_reject && ts->max_reject >= 0) {
354:       ts->reason = TS_DIVERGED_STEP_REJECTED;
355:       PetscInfo(ts,"Step=%D, step rejections %D greater than current TS allowed, stopping solve\n",ts->steps,rejections);
356:     }
357:   }
358:   return 0;
359: }

361: static PetscErrorCode TSInterpolate_IRK(TS ts,PetscReal itime,Vec U)
362: {
363:   TS_IRK          *irk = (TS_IRK*)ts->data;
364:   PetscInt        nstages = irk->nstages,pinterp = irk->pinterp,i,j;
365:   PetscReal       h;
366:   PetscReal       tt,t;
367:   PetscScalar     *bt;
368:   const PetscReal *B = irk->tableau->binterp;

371:   switch (irk->status) {
372:   case TS_STEP_INCOMPLETE:
373:   case TS_STEP_PENDING:
374:     h = ts->time_step;
375:     t = (itime - ts->ptime)/h;
376:     break;
377:   case TS_STEP_COMPLETE:
378:     h = ts->ptime - ts->ptime_prev;
379:     t = (itime - ts->ptime)/h + 1; /* In the interval [0,1] */
380:     break;
381:   default: SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_PLIB,"Invalid TSStepStatus");
382:   }
383:   PetscMalloc1(nstages,&bt);
384:   for (i=0; i<nstages; i++) bt[i] = 0;
385:   for (j=0,tt=t; j<pinterp; j++,tt*=t) {
386:     for (i=0; i<nstages; i++) {
387:       bt[i] += h * B[i*pinterp+j] * tt;
388:     }
389:   }
390:   VecMAXPY(U,nstages,bt,irk->YdotI);
391:   return 0;
392: }

394: static PetscErrorCode TSIRKTableauReset(TS ts)
395: {
396:   TS_IRK         *irk = (TS_IRK*)ts->data;
397:   IRKTableau     tab = irk->tableau;

399:   if (!tab) return 0;
400:   PetscFree3(tab->A,tab->A_inv,tab->I_s);
401:   PetscFree4(tab->b,tab->c,tab->binterp,tab->A_inv_rowsum);
402:   return 0;
403: }

405: static PetscErrorCode TSReset_IRK(TS ts)
406: {
407:   TS_IRK         *irk = (TS_IRK*)ts->data;

409:   TSIRKTableauReset(ts);
410:   if (irk->tableau) {
411:     PetscFree(irk->tableau);
412:   }
413:   if (irk->method_name) {
414:     PetscFree(irk->method_name);
415:   }
416:   if (irk->work) {
417:     PetscFree(irk->work);
418:   }
419:   VecDestroyVecs(irk->nstages,&irk->Y);
420:   VecDestroyVecs(irk->nstages,&irk->YdotI);
421:   VecDestroy(&irk->Ydot);
422:   VecDestroy(&irk->Z);
423:   VecDestroy(&irk->U);
424:   VecDestroy(&irk->U0);
425:   MatDestroy(&irk->TJ);
426:   return 0;
427: }

429: static PetscErrorCode TSIRKGetVecs(TS ts,DM dm,Vec *U)
430: {
431:   TS_IRK         *irk = (TS_IRK*)ts->data;

433:   if (U) {
434:     if (dm && dm != ts->dm) {
435:       DMGetNamedGlobalVector(dm,"TSIRK_U",U);
436:     } else *U = irk->U;
437:   }
438:   return 0;
439: }

441: static PetscErrorCode TSIRKRestoreVecs(TS ts,DM dm,Vec *U)
442: {
443:   if (U) {
444:     if (dm && dm != ts->dm) {
445:       DMRestoreNamedGlobalVector(dm,"TSIRK_U",U);
446:     }
447:   }
448:   return 0;
449: }

451: /*
452:   This defines the nonlinear equations that is to be solved with SNES
453:     G[e\otimes t + C*dt, Z, Zdot] = 0
454:     Zdot = (In \otimes S)*Z - (In \otimes Se) U
455:   where S = 1/(dt*A)
456: */
457: static PetscErrorCode SNESTSFormFunction_IRK(SNES snes,Vec ZC,Vec FC,TS ts)
458: {
459:   TS_IRK            *irk = (TS_IRK*)ts->data;
460:   IRKTableau        tab  = irk->tableau;
461:   const PetscInt    nstages = irk->nstages;
462:   const PetscReal   *c = tab->c;
463:   const PetscScalar *A_inv = tab->A_inv,*A_inv_rowsum = tab->A_inv_rowsum;
464:   DM                dm,dmsave;
465:   Vec               U,*YdotI = irk->YdotI,Ydot = irk->Ydot,*Y = irk->Y;
466:   PetscReal         h = ts->time_step;
467:   PetscInt          i,j;

469:   SNESGetDM(snes,&dm);
470:   TSIRKGetVecs(ts,dm,&U);
471:   VecStrideGatherAll(ZC,Y,INSERT_VALUES);
472:   dmsave = ts->dm;
473:   ts->dm = dm;
474:   for (i=0; i<nstages; i++) {
475:     VecZeroEntries(Ydot);
476:     for (j=0; j<nstages; j++) {
477:       VecAXPY(Ydot,A_inv[j*nstages+i]/h,Y[j]);
478:     }
479:     VecAXPY(Ydot,-A_inv_rowsum[i]/h,U); /* Ydot = (S \otimes In)*Z - (Se \otimes In) U */
480:     TSComputeIFunction(ts,ts->ptime+ts->time_step*c[i],Y[i],Ydot,YdotI[i],PETSC_FALSE);
481:   }
482:   VecStrideScatterAll(YdotI,FC,INSERT_VALUES);
483:   ts->dm = dmsave;
484:   TSIRKRestoreVecs(ts,dm,&U);
485:   return 0;
486: }

488: /*
489:    For explicit ODE, the Jacobian is
490:      JC = I_n \otimes S - J \otimes I_s
491:    For DAE, the Jacobian is
492:      JC = M_n \otimes S - J \otimes I_s
493: */
494: static PetscErrorCode SNESTSFormJacobian_IRK(SNES snes,Vec ZC,Mat JC,Mat JCpre,TS ts)
495: {
496:   TS_IRK          *irk = (TS_IRK*)ts->data;
497:   IRKTableau      tab  = irk->tableau;
498:   const PetscInt  nstages = irk->nstages;
499:   const PetscReal *c = tab->c;
500:   DM              dm,dmsave;
501:   Vec             *Y = irk->Y,Ydot = irk->Ydot;
502:   Mat             J;
503:   PetscScalar     *S;
504:   PetscInt        i,j,bs;

506:   SNESGetDM(snes,&dm);
507:   /* irk->Ydot has already been computed in SNESTSFormFunction_IRK (SNES guarantees this) */
508:   dmsave = ts->dm;
509:   ts->dm = dm;
510:   VecGetBlockSize(Y[nstages-1],&bs);
511:   if (ts->equation_type <= TS_EQ_ODE_EXPLICIT) { /* Support explicit formulas only */
512:     VecStrideGather(ZC,(nstages-1)*bs,Y[nstages-1],INSERT_VALUES);
513:     MatKAIJGetAIJ(JC,&J);
514:     TSComputeIJacobian(ts,ts->ptime+ts->time_step*c[nstages-1],Y[nstages-1],Ydot,0,J,J,PETSC_FALSE);
515:     MatKAIJGetS(JC,NULL,NULL,&S);
516:     for (i=0; i<nstages; i++)
517:       for (j=0; j<nstages; j++)
518:         S[i+nstages*j] = tab->A_inv[i+nstages*j]/ts->time_step;
519:     MatKAIJRestoreS(JC,&S);
520:   } else SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_SUP,"TSIRK %s does not support implicit formula",irk->method_name); /* TODO: need the mass matrix for DAE  */
521:   ts->dm = dmsave;
522:   return 0;
523: }

525: static PetscErrorCode DMCoarsenHook_TSIRK(DM fine,DM coarse,void *ctx)
526: {
527:   return 0;
528: }

530: static PetscErrorCode DMRestrictHook_TSIRK(DM fine,Mat restrct,Vec rscale,Mat inject,DM coarse,void *ctx)
531: {
532:   TS             ts = (TS)ctx;
533:   Vec            U,U_c;

535:   TSIRKGetVecs(ts,fine,&U);
536:   TSIRKGetVecs(ts,coarse,&U_c);
537:   MatRestrict(restrct,U,U_c);
538:   VecPointwiseMult(U_c,rscale,U_c);
539:   TSIRKRestoreVecs(ts,fine,&U);
540:   TSIRKRestoreVecs(ts,coarse,&U_c);
541:   return 0;
542: }

544: static PetscErrorCode DMSubDomainHook_TSIRK(DM dm,DM subdm,void *ctx)
545: {
546:   return 0;
547: }

549: static PetscErrorCode DMSubDomainRestrictHook_TSIRK(DM dm,VecScatter gscat,VecScatter lscat,DM subdm,void *ctx)
550: {
551:   TS             ts = (TS)ctx;
552:   Vec            U,U_c;

554:   TSIRKGetVecs(ts,dm,&U);
555:   TSIRKGetVecs(ts,subdm,&U_c);

557:   VecScatterBegin(gscat,U,U_c,INSERT_VALUES,SCATTER_FORWARD);
558:   VecScatterEnd(gscat,U,U_c,INSERT_VALUES,SCATTER_FORWARD);

560:   TSIRKRestoreVecs(ts,dm,&U);
561:   TSIRKRestoreVecs(ts,subdm,&U_c);
562:   return 0;
563: }

565: static PetscErrorCode TSSetUp_IRK(TS ts)
566: {
567:   TS_IRK         *irk = (TS_IRK*)ts->data;
568:   IRKTableau     tab = irk->tableau;
569:   DM             dm;
570:   Mat            J;
571:   Vec            R;
572:   const PetscInt nstages = irk->nstages;
573:   PetscInt       vsize,bs;

575:   if (!irk->work) {
576:     PetscMalloc1(irk->nstages,&irk->work);
577:   }
578:   if (!irk->Y) {
579:     VecDuplicateVecs(ts->vec_sol,irk->nstages,&irk->Y);
580:   }
581:   if (!irk->YdotI) {
582:     VecDuplicateVecs(ts->vec_sol,irk->nstages,&irk->YdotI);
583:   }
584:   if (!irk->Ydot) {
585:     VecDuplicate(ts->vec_sol,&irk->Ydot);
586:   }
587:   if (!irk->U) {
588:     VecDuplicate(ts->vec_sol,&irk->U);
589:   }
590:   if (!irk->U0) {
591:     VecDuplicate(ts->vec_sol,&irk->U0);
592:   }
593:   if (!irk->Z) {
594:     VecCreate(PetscObjectComm((PetscObject)ts->vec_sol),&irk->Z);
595:     VecGetSize(ts->vec_sol,&vsize);
596:     VecSetSizes(irk->Z,PETSC_DECIDE,vsize*irk->nstages);
597:     VecGetBlockSize(ts->vec_sol,&bs);
598:     VecSetBlockSize(irk->Z,irk->nstages*bs);
599:     VecSetFromOptions(irk->Z);
600:   }
601:   TSGetDM(ts,&dm);
602:   DMCoarsenHookAdd(dm,DMCoarsenHook_TSIRK,DMRestrictHook_TSIRK,ts);
603:   DMSubDomainHookAdd(dm,DMSubDomainHook_TSIRK,DMSubDomainRestrictHook_TSIRK,ts);

605:   TSGetSNES(ts,&ts->snes);
606:   VecDuplicate(irk->Z,&R);
607:   SNESSetFunction(ts->snes,R,SNESTSFormFunction,ts);
608:   TSGetIJacobian(ts,&J,NULL,NULL,NULL);
609:   if (!irk->TJ) {
610:     /* Create the KAIJ matrix for solving the stages */
611:     MatCreateKAIJ(J,nstages,nstages,tab->A_inv,tab->I_s,&irk->TJ);
612:   }
613:   SNESSetJacobian(ts->snes,irk->TJ,irk->TJ,SNESTSFormJacobian,ts);
614:   VecDestroy(&R);
615:   return 0;
616: }

618: static PetscErrorCode TSSetFromOptions_IRK(PetscOptionItems *PetscOptionsObject,TS ts)
619: {
620:   TS_IRK         *irk = (TS_IRK*)ts->data;
621:   char           tname[256] = TSIRKGAUSS;

623:   PetscOptionsHead(PetscOptionsObject,"IRK ODE solver options");
624:   {
625:     PetscBool flg1,flg2;
626:     PetscOptionsInt("-ts_irk_nstages","Stages of the IRK method","TSIRKSetNumStages",irk->nstages,&irk->nstages,&flg1);
627:     PetscOptionsFList("-ts_irk_type","Type of IRK method","TSIRKSetType",TSIRKList,irk->method_name[0] ? irk->method_name : tname,tname,sizeof(tname),&flg2);
628:     if (flg1 ||flg2 || !irk->method_name[0]) { /* Create the method tableau after nstages or method is set */
629:       TSIRKSetType(ts,tname);
630:     }
631:   }
632:   PetscOptionsTail();
633:   return 0;
634: }

636: static PetscErrorCode TSView_IRK(TS ts,PetscViewer viewer)
637: {
638:   TS_IRK         *irk = (TS_IRK*)ts->data;
639:   PetscBool      iascii;

641:   PetscObjectTypeCompare((PetscObject)viewer,PETSCVIEWERASCII,&iascii);
642:   if (iascii) {
643:     IRKTableau    tab = irk->tableau;
644:     TSIRKType irktype;
645:     char          buf[512];

647:     TSIRKGetType(ts,&irktype);
648:     PetscViewerASCIIPrintf(viewer,"  IRK type %s\n",irktype);
649:     PetscFormatRealArray(buf,sizeof(buf),"% 8.6f",irk->nstages,tab->c);
650:     PetscViewerASCIIPrintf(viewer,"  Abscissa       c = %s\n",buf);
651:     PetscViewerASCIIPrintf(viewer,"Stiffly accurate: %s\n",irk->stiffly_accurate ? "yes" : "no");
652:     PetscFormatRealArray(buf,sizeof(buf),"% 8.6f",PetscSqr(irk->nstages),tab->A);
653:     PetscViewerASCIIPrintf(viewer,"  A coefficients       A = %s\n",buf);
654:   }
655:   return 0;
656: }

658: static PetscErrorCode TSLoad_IRK(TS ts,PetscViewer viewer)
659: {
660:   SNES           snes;
661:   TSAdapt        adapt;

663:   TSGetAdapt(ts,&adapt);
664:   TSAdaptLoad(adapt,viewer);
665:   TSGetSNES(ts,&snes);
666:   SNESLoad(snes,viewer);
667:   /* function and Jacobian context for SNES when used with TS is always ts object */
668:   SNESSetFunction(snes,NULL,NULL,ts);
669:   SNESSetJacobian(snes,NULL,NULL,NULL,ts);
670:   return 0;
671: }

673: /*@C
674:   TSIRKSetType - Set the type of IRK scheme

676:   Logically collective

678:   Input Parameters:
679: +  ts - timestepping context
680: -  irktype - type of IRK scheme

682:   Options Database:
683: .  -ts_irk_type <gauss> - set irk type

685:   Level: intermediate

687: .seealso: TSIRKGetType(), TSIRK, TSIRKType, TSIRKGAUSS
688: @*/
689: PetscErrorCode TSIRKSetType(TS ts,TSIRKType irktype)
690: {
693:   PetscTryMethod(ts,"TSIRKSetType_C",(TS,TSIRKType),(ts,irktype));
694:   return 0;
695: }

697: /*@C
698:   TSIRKGetType - Get the type of IRK IMEX scheme

700:   Logically collective

702:   Input Parameter:
703: .  ts - timestepping context

705:   Output Parameter:
706: .  irktype - type of IRK-IMEX scheme

708:   Level: intermediate

710: .seealso: TSIRKGetType()
711: @*/
712: PetscErrorCode TSIRKGetType(TS ts,TSIRKType *irktype)
713: {
715:   PetscUseMethod(ts,"TSIRKGetType_C",(TS,TSIRKType*),(ts,irktype));
716:   return 0;
717: }

719: /*@C
720:   TSIRKSetNumStages - Set the number of stages of IRK scheme

722:   Logically collective

724:   Input Parameters:
725: +  ts - timestepping context
726: -  nstages - number of stages of IRK scheme

728:   Options Database:
729: .  -ts_irk_nstages <int> - set number of stages

731:   Level: intermediate

733: .seealso: TSIRKGetNumStages(), TSIRK
734: @*/
735: PetscErrorCode TSIRKSetNumStages(TS ts,PetscInt nstages)
736: {
738:   PetscTryMethod(ts,"TSIRKSetNumStages_C",(TS,PetscInt),(ts,nstages));
739:   return 0;
740: }

742: /*@C
743:   TSIRKGetNumStages - Get the number of stages of IRK scheme

745:   Logically collective

747:   Input Parameters:
748: +  ts - timestepping context
749: -  nstages - number of stages of IRK scheme

751:   Level: intermediate

753: .seealso: TSIRKSetNumStages(), TSIRK
754: @*/
755: PetscErrorCode TSIRKGetNumStages(TS ts,PetscInt *nstages)
756: {
759:   PetscTryMethod(ts,"TSIRKGetNumStages_C",(TS,PetscInt*),(ts,nstages));
760:   return 0;
761: }

763: static PetscErrorCode TSIRKGetType_IRK(TS ts,TSIRKType *irktype)
764: {
765:   TS_IRK *irk = (TS_IRK*)ts->data;

767:   *irktype = irk->method_name;
768:   return 0;
769: }

771: static PetscErrorCode TSIRKSetType_IRK(TS ts,TSIRKType irktype)
772: {
773:   TS_IRK         *irk = (TS_IRK*)ts->data;
774:   PetscErrorCode (*irkcreate)(TS);

776:   if (irk->method_name) {
777:     PetscFree(irk->method_name);
778:     TSIRKTableauReset(ts);
779:   }
780:   PetscFunctionListFind(TSIRKList,irktype,&irkcreate);
782:   (*irkcreate)(ts);
783:   PetscStrallocpy(irktype,&irk->method_name);
784:   return 0;
785: }

787: static PetscErrorCode TSIRKSetNumStages_IRK(TS ts,PetscInt nstages)
788: {
789:   TS_IRK *irk = (TS_IRK*)ts->data;

792:   irk->nstages = nstages;
793:   return 0;
794: }

796: static PetscErrorCode TSIRKGetNumStages_IRK(TS ts,PetscInt *nstages)
797: {
798:   TS_IRK *irk = (TS_IRK*)ts->data;

801:   *nstages = irk->nstages;
802:   return 0;
803: }

805: static PetscErrorCode TSDestroy_IRK(TS ts)
806: {
807:   TSReset_IRK(ts);
808:   if (ts->dm) {
809:     DMCoarsenHookRemove(ts->dm,DMCoarsenHook_TSIRK,DMRestrictHook_TSIRK,ts);
810:     DMSubDomainHookRemove(ts->dm,DMSubDomainHook_TSIRK,DMSubDomainRestrictHook_TSIRK,ts);
811:   }
812:   PetscFree(ts->data);
813:   PetscObjectComposeFunction((PetscObject)ts,"TSIRKSetType_C",NULL);
814:   PetscObjectComposeFunction((PetscObject)ts,"TSIRKGetType_C",NULL);
815:   PetscObjectComposeFunction((PetscObject)ts,"TSIRKSetNumStages_C",NULL);
816:   PetscObjectComposeFunction((PetscObject)ts,"TSIRKGetNumStages_C",NULL);
817:   return 0;
818: }

820: /*MC
821:       TSIRK - ODE and DAE solver using Implicit Runge-Kutta schemes

823:   Notes:

825:   TSIRK uses the sparse Kronecker product matrix implementation of MATKAIJ to achieve good arithmetic intensity.

827:   Gauss-Legrendre methods are currently supported. These are A-stable symplectic methods with an arbitrary number of stages. The order of accuracy is 2s when using s stages. The default method uses three stages and thus has an order of six. The number of stages (thus order) can be set with -ts_irk_nstages or TSIRKSetNumStages().

829:   Level: beginner

831: .seealso:  TSCreate(), TS, TSSetType(), TSIRKSetType(), TSIRKGetType(), TSIRKGAUSS, TSIRKRegister(), TSIRKSetNumStages()

833: M*/
834: PETSC_EXTERN PetscErrorCode TSCreate_IRK(TS ts)
835: {
836:   TS_IRK         *irk;

838:   TSIRKInitializePackage();

840:   ts->ops->reset          = TSReset_IRK;
841:   ts->ops->destroy        = TSDestroy_IRK;
842:   ts->ops->view           = TSView_IRK;
843:   ts->ops->load           = TSLoad_IRK;
844:   ts->ops->setup          = TSSetUp_IRK;
845:   ts->ops->step           = TSStep_IRK;
846:   ts->ops->interpolate    = TSInterpolate_IRK;
847:   ts->ops->evaluatestep   = TSEvaluateStep_IRK;
848:   ts->ops->rollback       = TSRollBack_IRK;
849:   ts->ops->setfromoptions = TSSetFromOptions_IRK;
850:   ts->ops->snesfunction   = SNESTSFormFunction_IRK;
851:   ts->ops->snesjacobian   = SNESTSFormJacobian_IRK;

853:   ts->usessnes = PETSC_TRUE;

855:   PetscNewLog(ts,&irk);
856:   ts->data = (void*)irk;

858:   PetscObjectComposeFunction((PetscObject)ts,"TSIRKSetType_C",TSIRKSetType_IRK);
859:   PetscObjectComposeFunction((PetscObject)ts,"TSIRKGetType_C",TSIRKGetType_IRK);
860:   PetscObjectComposeFunction((PetscObject)ts,"TSIRKSetNumStages_C",TSIRKSetNumStages_IRK);
861:   PetscObjectComposeFunction((PetscObject)ts,"TSIRKGetNumStages_C",TSIRKGetNumStages_IRK);
862:   /* 3-stage IRK_Gauss is the default */
863:   PetscNew(&irk->tableau);
864:   irk->nstages = 3;
865:   TSIRKSetType(ts,TSIRKDefault);
866:   return 0;
867: }