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;

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

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

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

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

133:   MatDestroy(&G0mat);
134:   MatDestroy(&G1mat);
135:   MatDestroy(&Amat);
136:   PetscFree3(b,G0,G1);

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

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

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

170: /*@C
171:    TSIRKRegister -  adds a TSIRK implementation

173:    Not Collective

175:    Input Parameters:
176: +  sname - name of user-defined IRK scheme
177: -  function - function to create method context

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

182:    Sample usage:
183: .vb
184:    TSIRKRegister("my_scheme",MySchemeCreate);
185: .ve

187:    Then, your scheme can be chosen with the procedural interface via
188: $     TSIRKSetType(ts,"my_scheme")
189:    or at runtime via the option
190: $     -ts_irk_type my_scheme

192:    Level: advanced

194: .seealso: TSIRKRegisterAll()
195: @*/
196: PetscErrorCode TSIRKRegister(const char sname[],PetscErrorCode (*function)(TS))
197: {

201:   TSIRKInitializePackage();
202:   PetscFunctionListAdd(&TSIRKList,sname,function);
203:   return(0);
204: }

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

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

211:   Level: advanced

213: .seealso:  TSIRKRegisterDestroy()
214: @*/
215: PetscErrorCode TSIRKRegisterAll(void)
216: {

220:   if (TSIRKRegisterAllCalled) return(0);
221:   TSIRKRegisterAllCalled = PETSC_TRUE;

223:   TSIRKRegister(TSIRKGAUSS,TSIRKCreate_Gauss);
224:   return(0);
225: }

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

230:    Not Collective

232:    Level: advanced

234: .seealso: TSIRKRegister(), TSIRKRegisterAll()
235: @*/
236: PetscErrorCode TSIRKRegisterDestroy(void)
237: {
239:   TSIRKRegisterAllCalled = PETSC_FALSE;
240:   return(0);
241: }

243: /*@C
244:   TSIRKInitializePackage - This function initializes everything in the TSIRK package. It is called
245:   from TSInitializePackage().

247:   Level: developer

249: .seealso: PetscInitialize()
250: @*/
251: PetscErrorCode TSIRKInitializePackage(void)
252: {

256:   if (TSIRKPackageInitialized) return(0);
257:   TSIRKPackageInitialized = PETSC_TRUE;
258:   TSIRKRegisterAll();
259:   PetscRegisterFinalize(TSIRKFinalizePackage);
260:   return(0);
261: }

263: /*@C
264:   TSIRKFinalizePackage - This function destroys everything in the TSIRK package. It is
265:   called from PetscFinalize().

267:   Level: developer

269: .seealso: PetscFinalize()
270: @*/
271: PetscErrorCode TSIRKFinalizePackage(void)
272: {

276:   PetscFunctionListDestroy(&TSIRKList);
277:   TSIRKPackageInitialized = PETSC_FALSE;
278:   return(0);
279: }

281: /*
282:  This function can be called before or after ts->vec_sol has been updated.
283: */
284: static PetscErrorCode TSEvaluateStep_IRK(TS ts,PetscInt order,Vec U,PetscBool *done)
285: {
286:   TS_IRK         *irk = (TS_IRK*)ts->data;
287:   IRKTableau     tab = irk->tableau;
288:   Vec            *YdotI = irk->YdotI;
289:   PetscScalar    *w = irk->work;
290:   PetscReal      h;
291:   PetscInt       j;

295:   switch (irk->status) {
296:   case TS_STEP_INCOMPLETE:
297:   case TS_STEP_PENDING:
298:     h = ts->time_step; break;
299:   case TS_STEP_COMPLETE:
300:     h = ts->ptime - ts->ptime_prev; break;
301:   default: SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_PLIB,"Invalid TSStepStatus");
302:   }

304:   VecCopy(ts->vec_sol,U);
305:   for (j=0; j<irk->nstages; j++) w[j] = h*tab->b[j];
306:   VecMAXPY(U,irk->nstages,w,YdotI);
307:   return(0);
308: }

310: static PetscErrorCode TSRollBack_IRK(TS ts)
311: {
312:   TS_IRK         *irk = (TS_IRK*)ts->data;

316:   VecCopy(irk->U0,ts->vec_sol);
317:   return(0);
318: }

320: static PetscErrorCode TSStep_IRK(TS ts)
321: {
322:   TS_IRK          *irk = (TS_IRK*)ts->data;
323:   IRKTableau      tab = irk->tableau;
324:   PetscScalar     *A_inv = tab->A_inv,*A_inv_rowsum = tab->A_inv_rowsum;
325:   const PetscInt  nstages = irk->nstages;
326:   SNES            snes;
327:   PetscInt        i,j,its,lits,bs;
328:   TSAdapt         adapt;
329:   PetscInt        rejections = 0;
330:   PetscBool       accept = PETSC_TRUE;
331:   PetscReal       next_time_step = ts->time_step;
332:   PetscErrorCode  ierr;

335:   if (!ts->steprollback) {
336:     VecCopy(ts->vec_sol,irk->U0);
337:   }
338:   VecGetBlockSize(ts->vec_sol,&bs);
339:   for (i=0; i<nstages; i++) {
340:     VecStrideScatter(ts->vec_sol,i*bs,irk->Z,INSERT_VALUES);
341:   }

343:   irk->status = TS_STEP_INCOMPLETE;
344:   while (!ts->reason && irk->status != TS_STEP_COMPLETE) {
345:     VecCopy(ts->vec_sol,irk->U);
346:     TSGetSNES(ts,&snes);
347:     SNESSolve(snes,NULL,irk->Z);
348:     SNESGetIterationNumber(snes,&its);
349:     SNESGetLinearSolveIterations(snes,&lits);
350:     ts->snes_its += its; ts->ksp_its += lits;
351:     VecStrideGatherAll(irk->Z,irk->Y,INSERT_VALUES);
352:     for (i=0; i<nstages; i++) {
353:       VecZeroEntries(irk->YdotI[i]);
354:       for (j=0; j<nstages; j++) {
355:         VecAXPY(irk->YdotI[i],A_inv[i+j*nstages]/ts->time_step,irk->Y[j]);
356:       }
357:       VecAXPY(irk->YdotI[i],-A_inv_rowsum[i]/ts->time_step,irk->U);
358:     }
359:     irk->status = TS_STEP_INCOMPLETE;
360:     TSEvaluateStep_IRK(ts,irk->order,ts->vec_sol,NULL);
361:     irk->status = TS_STEP_PENDING;
362:     TSGetAdapt(ts,&adapt);
363:     TSAdaptChoose(adapt,ts,ts->time_step,NULL,&next_time_step,&accept);
364:     irk->status = accept ? TS_STEP_COMPLETE : TS_STEP_INCOMPLETE;
365:     if (!accept) {
366:       TSRollBack_IRK(ts);
367:       ts->time_step = next_time_step;
368:       goto reject_step;
369:     }

371:     ts->ptime += ts->time_step;
372:     ts->time_step = next_time_step;
373:     break;
374:   reject_step:
375:     ts->reject++; accept = PETSC_FALSE;
376:     if (!ts->reason && ++rejections > ts->max_reject && ts->max_reject >= 0) {
377:       ts->reason = TS_DIVERGED_STEP_REJECTED;
378:       PetscInfo2(ts,"Step=%D, step rejections %D greater than current TS allowed, stopping solve\n",ts->steps,rejections);
379:     }
380:   }
381:   return(0);
382: }

384: static PetscErrorCode TSInterpolate_IRK(TS ts,PetscReal itime,Vec U)
385: {
386:   TS_IRK          *irk = (TS_IRK*)ts->data;
387:   PetscInt        nstages = irk->nstages,pinterp = irk->pinterp,i,j;
388:   PetscReal       h;
389:   PetscReal       tt,t;
390:   PetscScalar     *bt;
391:   const PetscReal *B = irk->tableau->binterp;
392:   PetscErrorCode  ierr;

395:   if (!B) SETERRQ1(PetscObjectComm((PetscObject)ts),PETSC_ERR_SUP,"TSIRK %s does not have an interpolation formula",irk->method_name);
396:   switch (irk->status) {
397:   case TS_STEP_INCOMPLETE:
398:   case TS_STEP_PENDING:
399:     h = ts->time_step;
400:     t = (itime - ts->ptime)/h;
401:     break;
402:   case TS_STEP_COMPLETE:
403:     h = ts->ptime - ts->ptime_prev;
404:     t = (itime - ts->ptime)/h + 1; /* In the interval [0,1] */
405:     break;
406:   default: SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_PLIB,"Invalid TSStepStatus");
407:   }
408:   PetscMalloc1(nstages,&bt);
409:   for (i=0; i<nstages; i++) bt[i] = 0;
410:   for (j=0,tt=t; j<pinterp; j++,tt*=t) {
411:     for (i=0; i<nstages; i++) {
412:       bt[i] += h * B[i*pinterp+j] * tt;
413:     }
414:   }
415:   VecMAXPY(U,nstages,bt,irk->YdotI);
416:   return(0);
417: }

419: static PetscErrorCode TSIRKTableauReset(TS ts)
420: {
421:   TS_IRK         *irk = (TS_IRK*)ts->data;
422:   IRKTableau     tab = irk->tableau;

426:   if (!tab) return(0);
427:   PetscFree3(tab->A,tab->A_inv,tab->I_s);
428:   PetscFree4(tab->b,tab->c,tab->binterp,tab->A_inv_rowsum);
429:   return(0);
430: }

432: static PetscErrorCode TSReset_IRK(TS ts)
433: {
434:   TS_IRK         *irk = (TS_IRK*)ts->data;

438:   TSIRKTableauReset(ts);
439:   if (irk->tableau) {
440:     PetscFree(irk->tableau);
441:   }
442:   if (irk->method_name) {
443:     PetscFree(irk->method_name);
444:   }
445:   if (irk->work) {
446:     PetscFree(irk->work);
447:   }
448:   VecDestroyVecs(irk->nstages,&irk->Y);
449:   VecDestroyVecs(irk->nstages,&irk->YdotI);
450:   VecDestroy(&irk->Ydot);
451:   VecDestroy(&irk->Z);
452:   VecDestroy(&irk->U);
453:   VecDestroy(&irk->U0);
454:   MatDestroy(&irk->TJ);
455:   return(0);
456: }

458: static PetscErrorCode TSIRKGetVecs(TS ts,DM dm,Vec *U)
459: {
460:   TS_IRK         *irk = (TS_IRK*)ts->data;

464:   if (U) {
465:     if (dm && dm != ts->dm) {
466:       DMGetNamedGlobalVector(dm,"TSIRK_U",U);
467:     } else *U = irk->U;
468:   }
469:   return(0);
470: }

472: static PetscErrorCode TSIRKRestoreVecs(TS ts,DM dm,Vec *U)
473: {

477:   if (U) {
478:     if (dm && dm != ts->dm) {
479:       DMRestoreNamedGlobalVector(dm,"TSIRK_U",U);
480:     }
481:   }
482:   return(0);
483: }

485: /*
486:   This defines the nonlinear equations that is to be solved with SNES
487:     G[e\otimes t + C*dt, Z, Zdot] = 0
488:     Zdot = (In \otimes S)*Z - (In \otimes Se) U
489:   where S = 1/(dt*A)
490: */
491: static PetscErrorCode SNESTSFormFunction_IRK(SNES snes,Vec ZC,Vec FC,TS ts)
492: {
493:   TS_IRK            *irk = (TS_IRK*)ts->data;
494:   IRKTableau        tab  = irk->tableau;
495:   const PetscInt    nstages = irk->nstages;
496:   const PetscReal   *c = tab->c;
497:   const PetscScalar *A_inv = tab->A_inv,*A_inv_rowsum = tab->A_inv_rowsum;
498:   DM                dm,dmsave;
499:   Vec               U,*YdotI = irk->YdotI,Ydot = irk->Ydot,*Y = irk->Y;
500:   PetscReal         h = ts->time_step;
501:   PetscInt          i,j;
502:   PetscErrorCode    ierr;

505:   SNESGetDM(snes,&dm);
506:   TSIRKGetVecs(ts,dm,&U);
507:   VecStrideGatherAll(ZC,Y,INSERT_VALUES);
508:   dmsave = ts->dm;
509:   ts->dm = dm;
510:   for (i=0; i<nstages; i++) {
511:     VecZeroEntries(Ydot);
512:     for (j=0; j<nstages; j++) {
513:       VecAXPY(Ydot,A_inv[j*nstages+i]/h,Y[j]);
514:     }
515:     VecAXPY(Ydot,-A_inv_rowsum[i]/h,U); /* Ydot = (S \otimes In)*Z - (Se \otimes In) U */
516:     TSComputeIFunction(ts,ts->ptime+ts->time_step*c[i],Y[i],Ydot,YdotI[i],PETSC_FALSE);
517:   }
518:   VecStrideScatterAll(YdotI,FC,INSERT_VALUES);
519:   ts->dm = dmsave;
520:   TSIRKRestoreVecs(ts,dm,&U);
521:   return(0);
522: }

524: /*
525:    For explicit ODE, the Jacobian is
526:      JC = I_n \otimes S - J \otimes I_s
527:    For DAE, the Jacobian is
528:      JC = M_n \otimes S - J \otimes I_s
529: */
530: static PetscErrorCode SNESTSFormJacobian_IRK(SNES snes,Vec ZC,Mat JC,Mat JCpre,TS ts)
531: {
532:   TS_IRK          *irk = (TS_IRK*)ts->data;
533:   IRKTableau      tab  = irk->tableau;
534:   const PetscInt  nstages = irk->nstages;
535:   const PetscReal *c = tab->c;
536:   DM              dm,dmsave;
537:   Vec             *Y = irk->Y,Ydot = irk->Ydot;
538:   Mat             J;
539:   PetscScalar     *S;
540:   PetscInt        i,j,bs;
541:   PetscErrorCode  ierr;

544:   SNESGetDM(snes,&dm);
545:   /* irk->Ydot has already been computed in SNESTSFormFunction_IRK (SNES guarantees this) */
546:   dmsave = ts->dm;
547:   ts->dm = dm;
548:   VecGetBlockSize(Y[nstages-1],&bs);
549:   if (ts->equation_type <= TS_EQ_ODE_EXPLICIT) { /* Support explicit formulas only */
550:     VecStrideGather(ZC,(nstages-1)*bs,Y[nstages-1],INSERT_VALUES);
551:     MatKAIJGetAIJ(JC,&J);
552:     TSComputeIJacobian(ts,ts->ptime+ts->time_step*c[nstages-1],Y[nstages-1],Ydot,0,J,J,PETSC_FALSE);
553:     MatKAIJGetS(JC,NULL,NULL,&S);
554:     for (i=0; i<nstages; i++)
555:       for (j=0; j<nstages; j++)
556:         S[i+nstages*j] = tab->A_inv[i+nstages*j]/ts->time_step;
557:     MatKAIJRestoreS(JC,&S);
558:   } else
559:     SETERRQ1(PetscObjectComm((PetscObject)ts),PETSC_ERR_SUP,"TSIRK %s does not support implicit formula",irk->method_name); /* ToDo: need the mass matrix for DAE  */
560:   ts->dm = dmsave;
561:   return(0);
562: }

564: static PetscErrorCode DMCoarsenHook_TSIRK(DM fine,DM coarse,void *ctx)
565: {
567:   return(0);
568: }

570: static PetscErrorCode DMRestrictHook_TSIRK(DM fine,Mat restrct,Vec rscale,Mat inject,DM coarse,void *ctx)
571: {
572:   TS             ts = (TS)ctx;
573:   Vec            U,U_c;

577:   TSIRKGetVecs(ts,fine,&U);
578:   TSIRKGetVecs(ts,coarse,&U_c);
579:   MatRestrict(restrct,U,U_c);
580:   VecPointwiseMult(U_c,rscale,U_c);
581:   TSIRKRestoreVecs(ts,fine,&U);
582:   TSIRKRestoreVecs(ts,coarse,&U_c);
583:   return(0);
584: }

586: static PetscErrorCode DMSubDomainHook_TSIRK(DM dm,DM subdm,void *ctx)
587: {
589:   return(0);
590: }

592: static PetscErrorCode DMSubDomainRestrictHook_TSIRK(DM dm,VecScatter gscat,VecScatter lscat,DM subdm,void *ctx)
593: {
594:   TS             ts = (TS)ctx;
595:   Vec            U,U_c;

599:   TSIRKGetVecs(ts,dm,&U);
600:   TSIRKGetVecs(ts,subdm,&U_c);

602:   VecScatterBegin(gscat,U,U_c,INSERT_VALUES,SCATTER_FORWARD);
603:   VecScatterEnd(gscat,U,U_c,INSERT_VALUES,SCATTER_FORWARD);

605:   TSIRKRestoreVecs(ts,dm,&U);
606:   TSIRKRestoreVecs(ts,subdm,&U_c);
607:   return(0);
608: }

610: static PetscErrorCode TSSetUp_IRK(TS ts)
611: {
612:   TS_IRK         *irk = (TS_IRK*)ts->data;
613:   IRKTableau     tab = irk->tableau;
614:   DM             dm;
615:   Mat            J;
616:   Vec            R;
617:   const PetscInt nstages = irk->nstages;
618:   PetscInt       vsize,bs;

622:   if (!irk->work) {
623:     PetscMalloc1(irk->nstages,&irk->work);
624:   }
625:   if (!irk->Y) {
626:     VecDuplicateVecs(ts->vec_sol,irk->nstages,&irk->Y);
627:   }
628:   if (!irk->YdotI) {
629:     VecDuplicateVecs(ts->vec_sol,irk->nstages,&irk->YdotI);
630:   }
631:   if (!irk->Ydot) {
632:     VecDuplicate(ts->vec_sol,&irk->Ydot);
633:   }
634:   if (!irk->U) {
635:     VecDuplicate(ts->vec_sol,&irk->U);
636:   }
637:   if (!irk->U0) {
638:     VecDuplicate(ts->vec_sol,&irk->U0);
639:   }
640:   if (!irk->Z) {
641:     VecCreate(PetscObjectComm((PetscObject)ts->vec_sol),&irk->Z);
642:     VecGetSize(ts->vec_sol,&vsize);
643:     VecSetSizes(irk->Z,PETSC_DECIDE,vsize*irk->nstages);
644:     VecGetBlockSize(ts->vec_sol,&bs);
645:     VecSetBlockSize(irk->Z,irk->nstages*bs);
646:     VecSetFromOptions(irk->Z);
647:   }
648:   TSGetDM(ts,&dm);
649:   DMCoarsenHookAdd(dm,DMCoarsenHook_TSIRK,DMRestrictHook_TSIRK,ts);
650:   DMSubDomainHookAdd(dm,DMSubDomainHook_TSIRK,DMSubDomainRestrictHook_TSIRK,ts);

652:   TSGetSNES(ts,&ts->snes);
653:   VecDuplicate(irk->Z,&R);
654:   SNESSetFunction(ts->snes,R,SNESTSFormFunction,ts);
655:   TSGetIJacobian(ts,&J,NULL,NULL,NULL);
656:   if (!irk->TJ) {
657:     /* Create the KAIJ matrix for solving the stages */
658:     MatCreateKAIJ(J,nstages,nstages,tab->A_inv,tab->I_s,&irk->TJ);
659:   }
660:   SNESSetJacobian(ts->snes,irk->TJ,irk->TJ,SNESTSFormJacobian,ts);
661:   VecDestroy(&R);
662:   return(0);
663: }

665: static PetscErrorCode TSSetFromOptions_IRK(PetscOptionItems *PetscOptionsObject,TS ts)
666: {
667:   TS_IRK         *irk = (TS_IRK*)ts->data;
668:   char           tname[256] = TSIRKGAUSS;

672:   PetscOptionsHead(PetscOptionsObject,"IRK ODE solver options");
673:   {
674:     PetscBool flg1,flg2;
675:     PetscOptionsInt("-ts_irk_nstages","Stages of the IRK method","TSIRKSetNumStages",irk->nstages,&irk->nstages,&flg1);
676:     PetscOptionsFList("-ts_irk_type","Type of IRK method","TSIRKSetType",TSIRKList,irk->method_name[0] ? irk->method_name : tname,tname,sizeof(tname),&flg2);
677:     if (flg1 ||flg2 || !irk->method_name[0]) { /* Create the method tableau after nstages or method is set */
678:       TSIRKSetType(ts,tname);
679:     }
680:   }
681:   PetscOptionsTail();
682:   return(0);
683: }

685: static PetscErrorCode TSView_IRK(TS ts,PetscViewer viewer)
686: {
687:   TS_IRK         *irk = (TS_IRK*)ts->data;
688:   PetscBool      iascii;

692:   PetscObjectTypeCompare((PetscObject)viewer,PETSCVIEWERASCII,&iascii);
693:   if (iascii) {
694:     IRKTableau    tab = irk->tableau;
695:     TSIRKType irktype;
696:     char          buf[512];

698:     TSIRKGetType(ts,&irktype);
699:     PetscViewerASCIIPrintf(viewer,"  IRK type %s\n",irktype);
700:     PetscFormatRealArray(buf,sizeof(buf),"% 8.6f",irk->nstages,tab->c);
701:     PetscViewerASCIIPrintf(viewer,"  Abscissa       c = %s\n",buf);
702:     PetscViewerASCIIPrintf(viewer,"Stiffly accurate: %s\n",irk->stiffly_accurate ? "yes" : "no");
703:     PetscFormatRealArray(buf,sizeof(buf),"% 8.6f",PetscSqr(irk->nstages),tab->A);
704:     PetscViewerASCIIPrintf(viewer,"  A coefficients       A = %s\n",buf);
705:   }
706:   return(0);
707: }

709: static PetscErrorCode TSLoad_IRK(TS ts,PetscViewer viewer)
710: {
712:   SNES           snes;
713:   TSAdapt        adapt;

716:   TSGetAdapt(ts,&adapt);
717:   TSAdaptLoad(adapt,viewer);
718:   TSGetSNES(ts,&snes);
719:   SNESLoad(snes,viewer);
720:   /* function and Jacobian context for SNES when used with TS is always ts object */
721:   SNESSetFunction(snes,NULL,NULL,ts);
722:   SNESSetJacobian(snes,NULL,NULL,NULL,ts);
723:   return(0);
724: }

726: /*@C
727:   TSIRKSetType - Set the type of IRK scheme

729:   Logically collective

731:   Input Parameters:
732: +  ts - timestepping context
733: -  irktype - type of IRK scheme

735:   Options Database:
736: .  -ts_irk_type <gauss>

738:   Level: intermediate

740: .seealso: TSIRKGetType(), TSIRK, TSIRKType, TSIRKGAUSS
741: @*/
742: PetscErrorCode TSIRKSetType(TS ts,TSIRKType irktype)
743: {

749:   PetscTryMethod(ts,"TSIRKSetType_C",(TS,TSIRKType),(ts,irktype));
750:   return(0);
751: }

753: /*@C
754:   TSIRKGetType - Get the type of IRK IMEX scheme

756:   Logically collective

758:   Input Parameter:
759: .  ts - timestepping context

761:   Output Parameter:
762: .  irktype - type of IRK-IMEX scheme

764:   Level: intermediate

766: .seealso: TSIRKGetType()
767: @*/
768: PetscErrorCode TSIRKGetType(TS ts,TSIRKType *irktype)
769: {

774:   PetscUseMethod(ts,"TSIRKGetType_C",(TS,TSIRKType*),(ts,irktype));
775:   return(0);
776: }

778: /*@C
779:   TSIRKSetNumStages - Set the number of stages of IRK scheme

781:   Logically collective

783:   Input Parameters:
784: +  ts - timestepping context
785: -  nstages - number of stages of IRK scheme

787:   Options Database:
788: .  -ts_irk_nstages <int>: Number of stages

790:   Level: intermediate

792: .seealso: TSIRKGetNumStages(), TSIRK
793: @*/
794: PetscErrorCode TSIRKSetNumStages(TS ts,PetscInt nstages)
795: {

800:   PetscTryMethod(ts,"TSIRKSetNumStages_C",(TS,PetscInt),(ts,nstages));
801:   return(0);
802: }

804: /*@C
805:   TSIRKGetNumStages - Get the number of stages of IRK scheme

807:   Logically collective

809:   Input Parameters:
810: +  ts - timestepping context
811: -  nstages - number of stages of IRK scheme

813:   Level: intermediate

815: .seealso: TSIRKSetNumStages(), TSIRK
816: @*/
817: PetscErrorCode TSIRKGetNumStages(TS ts,PetscInt *nstages)
818: {

824:   PetscTryMethod(ts,"TSIRKGetNumStages_C",(TS,PetscInt*),(ts,nstages));
825:   return(0);
826: }

828: static PetscErrorCode TSIRKGetType_IRK(TS ts,TSIRKType *irktype)
829: {
830:   TS_IRK *irk = (TS_IRK*)ts->data;

833:   *irktype = irk->method_name;
834:   return(0);
835: }

837: static PetscErrorCode TSIRKSetType_IRK(TS ts,TSIRKType irktype)
838: {
839:   TS_IRK         *irk = (TS_IRK*)ts->data;
840:   PetscErrorCode ierr,(*irkcreate)(TS);

843:   if (irk->method_name) {
844:     PetscFree(irk->method_name);
845:     TSIRKTableauReset(ts);
846:   }
847:   PetscFunctionListFind(TSIRKList,irktype,&irkcreate);
848:   if (!irkcreate) SETERRQ1(PETSC_COMM_SELF,PETSC_ERR_ARG_UNKNOWN_TYPE,"Unknown TSIRK type \"%s\" given",irktype);
849:   (*irkcreate)(ts);
850:   PetscStrallocpy(irktype,&irk->method_name);
851:   return(0);
852: }

854: static PetscErrorCode TSIRKSetNumStages_IRK(TS ts,PetscInt nstages)
855: {
856:   TS_IRK *irk = (TS_IRK*)ts->data;

859:   if (nstages<=0) SETERRQ1(PETSC_COMM_SELF,PETSC_ERR_ARG_OUTOFRANGE,"input argument, %d, out of range",nstages);
860:   irk->nstages = nstages;
861:   return(0);
862: }

864: static PetscErrorCode TSIRKGetNumStages_IRK(TS ts,PetscInt *nstages)
865: {
866:   TS_IRK *irk = (TS_IRK*)ts->data;

870:   *nstages = irk->nstages;
871:   return(0);
872: }

874: static PetscErrorCode TSDestroy_IRK(TS ts)
875: {

879:   TSReset_IRK(ts);
880:   if (ts->dm) {
881:     DMCoarsenHookRemove(ts->dm,DMCoarsenHook_TSIRK,DMRestrictHook_TSIRK,ts);
882:     DMSubDomainHookRemove(ts->dm,DMSubDomainHook_TSIRK,DMSubDomainRestrictHook_TSIRK,ts);
883:   }
884:   PetscFree(ts->data);
885:   PetscObjectComposeFunction((PetscObject)ts,"TSIRKSetType_C",NULL);
886:   PetscObjectComposeFunction((PetscObject)ts,"TSIRKGetType_C",NULL);
887:   PetscObjectComposeFunction((PetscObject)ts,"TSIRKSetNumStages_C",NULL);
888:   PetscObjectComposeFunction((PetscObject)ts,"TSIRKGetNumStages_C",NULL);
889:   return(0);
890: }

892: /*MC
893:       TSIRK - ODE and DAE solver using Implicit Runge-Kutta schemes

895:   Notes:

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

899:   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().

901:   Level: beginner

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

905: M*/
906: PETSC_EXTERN PetscErrorCode TSCreate_IRK(TS ts)
907: {
908:   TS_IRK         *irk;

912:   TSIRKInitializePackage();

914:   ts->ops->reset          = TSReset_IRK;
915:   ts->ops->destroy        = TSDestroy_IRK;
916:   ts->ops->view           = TSView_IRK;
917:   ts->ops->load           = TSLoad_IRK;
918:   ts->ops->setup          = TSSetUp_IRK;
919:   ts->ops->step           = TSStep_IRK;
920:   ts->ops->interpolate    = TSInterpolate_IRK;
921:   ts->ops->evaluatestep   = TSEvaluateStep_IRK;
922:   ts->ops->rollback       = TSRollBack_IRK;
923:   ts->ops->setfromoptions = TSSetFromOptions_IRK;
924:   ts->ops->snesfunction   = SNESTSFormFunction_IRK;
925:   ts->ops->snesjacobian   = SNESTSFormJacobian_IRK;

927:   ts->usessnes = PETSC_TRUE;

929:   PetscNewLog(ts,&irk);
930:   ts->data = (void*)irk;

932:   PetscObjectComposeFunction((PetscObject)ts,"TSIRKSetType_C",TSIRKSetType_IRK);
933:   PetscObjectComposeFunction((PetscObject)ts,"TSIRKGetType_C",TSIRKGetType_IRK);
934:   PetscObjectComposeFunction((PetscObject)ts,"TSIRKSetNumStages_C",TSIRKSetNumStages_IRK);
935:   PetscObjectComposeFunction((PetscObject)ts,"TSIRKGetNumStages_C",TSIRKGetNumStages_IRK);
936:   /* 3-stage IRK_Gauss is the default */
937:   PetscNew(&irk->tableau);
938:   irk->nstages = 3;
939:   TSIRKSetType(ts,TSIRKDefault);
940:   return(0);
941: }