Actual source code: theta.c

petsc-master 2017-07-18
Report Typos and Errors
  1:  /*
  2:    Code for timestepping with implicit Theta method
  3:  */
  4:  #include <petsc/private/tsimpl.h>                /*I   "petscts.h"   I*/
  5:  #include <petscsnes.h>
  6:  #include <petscdm.h>
  7:  #include <petscmat.h>

  9:  typedef struct {
 10:    PetscReal    stage_time;
 11:    Vec          X0,X,Xdot;                /* Storage for stages and time derivative */
 12:    Vec          affine;                   /* Affine vector needed for residual at beginning of step in endpoint formulation */

 14:    PetscReal    Theta;
 15:    PetscInt     order;
 16:    PetscBool    endpoint;
 17:    PetscBool    extrapolate;

 19:    Vec          *VecsDeltaLam;            /* Increment of the adjoint sensitivity w.r.t IC at stage */
 20:    Vec          *VecsDeltaMu;             /* Increment of the adjoint sensitivity w.r.t P at stage */
 21:    Vec          *VecsSensiTemp;           /* Vector to be timed with Jacobian transpose */
 22:    Vec          VecCostIntegral0;         /* Backup for roll-backs due to events */
 23:    PetscReal    ptime;
 24:    PetscReal    time_step;

 26:    Vec          vec_sol_prev;
 27:    Vec          vec_lte_work;

 29:    TSStepStatus status;
 30:  } TS_Theta;

 32:  static PetscErrorCode TSThetaGetX0AndXdot(TS ts,DM dm,Vec *X0,Vec *Xdot)
 33:  {
 34:    TS_Theta       *th = (TS_Theta*)ts->data;

 38:    if (X0) {
 39:      if (dm && dm != ts->dm) {
 40:        DMGetNamedGlobalVector(dm,"TSTheta_X0",X0);
 41:      } else *X0 = ts->vec_sol;
 42:    }
 43:    if (Xdot) {
 44:      if (dm && dm != ts->dm) {
 45:        DMGetNamedGlobalVector(dm,"TSTheta_Xdot",Xdot);
 46:      } else *Xdot = th->Xdot;
 47:    }
 48:    return(0);
 49:  }

 51:  static PetscErrorCode TSThetaRestoreX0AndXdot(TS ts,DM dm,Vec *X0,Vec *Xdot)
 52:  {

 56:    if (X0) {
 57:      if (dm && dm != ts->dm) {
 58:        DMRestoreNamedGlobalVector(dm,"TSTheta_X0",X0);
 59:      }
 60:    }
 61:    if (Xdot) {
 62:      if (dm && dm != ts->dm) {
 63:        DMRestoreNamedGlobalVector(dm,"TSTheta_Xdot",Xdot);
 64:      }
 65:    }
 66:    return(0);
 67:  }

 69:  static PetscErrorCode DMCoarsenHook_TSTheta(DM fine,DM coarse,void *ctx)
 70:  {

 73:    return(0);
 74:  }

 76:  static PetscErrorCode DMRestrictHook_TSTheta(DM fine,Mat restrct,Vec rscale,Mat inject,DM coarse,void *ctx)
 77:  {
 78:    TS             ts = (TS)ctx;
 80:    Vec            X0,Xdot,X0_c,Xdot_c;

 83:    TSThetaGetX0AndXdot(ts,fine,&X0,&Xdot);
 84:    TSThetaGetX0AndXdot(ts,coarse,&X0_c,&Xdot_c);
 85:    MatRestrict(restrct,X0,X0_c);
 86:    MatRestrict(restrct,Xdot,Xdot_c);
 87:    VecPointwiseMult(X0_c,rscale,X0_c);
 88:    VecPointwiseMult(Xdot_c,rscale,Xdot_c);
 89:    TSThetaRestoreX0AndXdot(ts,fine,&X0,&Xdot);
 90:    TSThetaRestoreX0AndXdot(ts,coarse,&X0_c,&Xdot_c);
 91:    return(0);
 92:  }

 94:  static PetscErrorCode DMSubDomainHook_TSTheta(DM dm,DM subdm,void *ctx)
 95:  {

 98:    return(0);
 99:  }

101:  static PetscErrorCode DMSubDomainRestrictHook_TSTheta(DM dm,VecScatter gscat,VecScatter lscat,DM subdm,void *ctx)
102:  {
103:    TS             ts = (TS)ctx;
105:    Vec            X0,Xdot,X0_sub,Xdot_sub;

108:    TSThetaGetX0AndXdot(ts,dm,&X0,&Xdot);
109:    TSThetaGetX0AndXdot(ts,subdm,&X0_sub,&Xdot_sub);

111:    VecScatterBegin(gscat,X0,X0_sub,INSERT_VALUES,SCATTER_FORWARD);
112:    VecScatterEnd(gscat,X0,X0_sub,INSERT_VALUES,SCATTER_FORWARD);

114:    VecScatterBegin(gscat,Xdot,Xdot_sub,INSERT_VALUES,SCATTER_FORWARD);
115:    VecScatterEnd(gscat,Xdot,Xdot_sub,INSERT_VALUES,SCATTER_FORWARD);

117:    TSThetaRestoreX0AndXdot(ts,dm,&X0,&Xdot);
118:    TSThetaRestoreX0AndXdot(ts,subdm,&X0_sub,&Xdot_sub);
119:    return(0);
120:  }

122:  static PetscErrorCode TSForwardCostIntegral_Theta(TS ts)
123:  {
124:    TS_Theta       *th = (TS_Theta*)ts->data;

128:    /* backup cost integral */
129:    VecCopy(ts->vec_costintegral,th->VecCostIntegral0);
130:    if (th->endpoint) {
131:      TSAdjointComputeCostIntegrand(ts,th->ptime,th->X0,ts->vec_costintegrand);
132:      VecAXPY(ts->vec_costintegral,th->time_step*(1-th->Theta),ts->vec_costintegrand);
133:    }
134:    TSAdjointComputeCostIntegrand(ts,th->stage_time,th->X,ts->vec_costintegrand);
135:    if (th->endpoint) {
136:      VecAXPY(ts->vec_costintegral,th->time_step*th->Theta,ts->vec_costintegrand);
137:    } else {
138:      VecAXPY(ts->vec_costintegral,th->time_step,ts->vec_costintegrand);
139:    }
140:    return(0);
141:  }

143:  static PetscErrorCode TSAdjointCostIntegral_Theta(TS ts)
144:  {
145:    TS_Theta       *th = (TS_Theta*)ts->data;

149:    if (th->endpoint) {
150:      /* Evolve ts->vec_costintegral to compute integrals */
151:      TSAdjointComputeCostIntegrand(ts,ts->ptime,ts->vec_sol,ts->vec_costintegrand);
152:      VecAXPY(ts->vec_costintegral,-ts->time_step*th->Theta,ts->vec_costintegrand);
153:      if (th->Theta!=1) {
154:        TSAdjointComputeCostIntegrand(ts,th->ptime,th->X0,ts->vec_costintegrand);
155:        VecAXPY(ts->vec_costintegral,ts->time_step*(th->Theta-1),ts->vec_costintegrand);
156:      }
157:    }else {
158:      TSAdjointComputeCostIntegrand(ts,th->stage_time,th->X,ts->vec_costintegrand);
159:      VecAXPY(ts->vec_costintegral,-ts->time_step,ts->vec_costintegrand);
160:    }
161:    return(0);
162:  }

164:  static PetscErrorCode TS_SNESSolve(TS ts,Vec b,Vec x)
165:  {
166:    PetscInt       nits,lits;

170:    SNESSolve(ts->snes,b,x);
171:    SNESGetIterationNumber(ts->snes,&nits);
172:    SNESGetLinearSolveIterations(ts->snes,&lits);
173:    ts->snes_its += nits; ts->ksp_its += lits;
174:    return(0);
175:  }

177:  static PetscErrorCode TSStep_Theta(TS ts)
178:  {
179:    TS_Theta       *th = (TS_Theta*)ts->data;
180:    PetscInt       rejections = 0;
181:    PetscBool      stageok,accept = PETSC_TRUE;
182:    PetscReal      next_time_step = ts->time_step;

186:    if (!ts->steprollback) {
187:      if (th->vec_sol_prev) { VecCopy(th->X0,th->vec_sol_prev); }
188:      VecCopy(ts->vec_sol,th->X0);
189:    }

191:    th->status = TS_STEP_INCOMPLETE;
192:    while (!ts->reason && th->status != TS_STEP_COMPLETE) {

194:      PetscReal shift = 1/(th->Theta*ts->time_step);
195:      th->stage_time = ts->ptime + (th->endpoint ? (PetscReal)1 : th->Theta)*ts->time_step;

197:      VecCopy(th->X0,th->X);
198:      if (th->extrapolate && !ts->steprestart) {
199:        VecAXPY(th->X,1/shift,th->Xdot);
200:      }
201:      if (th->endpoint) { /* This formulation assumes linear time-independent mass matrix */
202:        if (!th->affine) {VecDuplicate(ts->vec_sol,&th->affine);}
203:        VecZeroEntries(th->Xdot);
204:        TSComputeIFunction(ts,ts->ptime,th->X0,th->Xdot,th->affine,PETSC_FALSE);
205:        VecScale(th->affine,(th->Theta-1)/th->Theta);
206:      } else if (th->affine) { /* Just in case th->endpoint is changed between calls to TSStep_Theta() */
207:        VecZeroEntries(th->affine);
208:      }
209:      TSPreStage(ts,th->stage_time);
210:      TS_SNESSolve(ts,th->affine,th->X);
211:      TSPostStage(ts,th->stage_time,0,&th->X);
212:      TSAdaptCheckStage(ts->adapt,ts,th->stage_time,th->X,&stageok);
213:      if (!stageok) goto reject_step;

215:      th->status = TS_STEP_PENDING;
216:      if (th->endpoint) {
217:        VecCopy(th->X,ts->vec_sol);
218:      } else {
219:        VecAXPBYPCZ(th->Xdot,-shift,shift,0,th->X0,th->X);
220:        VecAXPY(ts->vec_sol,ts->time_step,th->Xdot);
221:      }
222:      TSAdaptChoose(ts->adapt,ts,ts->time_step,NULL,&next_time_step,&accept);
223:      th->status = accept ? TS_STEP_COMPLETE : TS_STEP_INCOMPLETE;
224:      if (!accept) {
225:        VecCopy(th->X0,ts->vec_sol);
226:        ts->time_step = next_time_step;
227:        goto reject_step;
228:      }

230:      if (ts->costintegralfwd) { /* Save the info for the later use in cost integral evaluation */
231:        th->ptime     = ts->ptime;
232:        th->time_step = ts->time_step;
233:      }

235:      ts->ptime += ts->time_step;
236:      ts->time_step = next_time_step;
237:      break;

239:    reject_step:
240:      ts->reject++; accept = PETSC_FALSE;
241:      if (!ts->reason && ++rejections > ts->max_reject && ts->max_reject >= 0) {
242:        ts->reason = TS_DIVERGED_STEP_REJECTED;
243:        PetscInfo2(ts,"Step=%D, step rejections %D greater than current TS allowed, stopping solve\n",ts->steps,rejections);
244:      }
245:    }
246:    return(0);
247:  }

249:  static PetscErrorCode TSAdjointStep_Theta(TS ts)
250:  {
251:    TS_Theta            *th = (TS_Theta*)ts->data;
252:    Vec                 *VecsDeltaLam = th->VecsDeltaLam,*VecsDeltaMu = th->VecsDeltaMu,*VecsSensiTemp = th->VecsSensiTemp;
253:    PetscInt            nadj;
254:    PetscErrorCode      ierr;
255:    Mat                 J,Jp;
256:    KSP                 ksp;
257:    PetscReal           shift;


261:    th->status = TS_STEP_INCOMPLETE;
262:    SNESGetKSP(ts->snes,&ksp);
263:    TSGetIJacobian(ts,&J,&Jp,NULL,NULL);

265:    /* If endpoint=1, th->ptime and th->X0 will be used; if endpoint=0, th->stage_time and th->X will be used. */
266:    th->stage_time = ts->ptime + (th->endpoint ? ts->time_step : (1.-th->Theta)*ts->time_step); /* time_step is negative*/
267:    th->ptime      = ts->ptime + ts->time_step;

269:    /* Build RHS */
270:    if (ts->vec_costintegral) { /* Cost function has an integral term */
271:      if (th->endpoint) {
272:        TSAdjointComputeDRDYFunction(ts,ts->ptime,ts->vec_sol,ts->vecs_drdy);
273:      }else {
274:        TSAdjointComputeDRDYFunction(ts,th->stage_time,th->X,ts->vecs_drdy);
275:      }
276:    }
277:    for (nadj=0; nadj<ts->numcost; nadj++) {
278:      VecCopy(ts->vecs_sensi[nadj],VecsSensiTemp[nadj]);
279:      VecScale(VecsSensiTemp[nadj],-1./(th->Theta*ts->time_step));
280:      if (ts->vec_costintegral) {
281:        VecAXPY(VecsSensiTemp[nadj],1.,ts->vecs_drdy[nadj]);
282:      }
283:    }

285:    /* Build LHS */
286:    shift = -1./(th->Theta*ts->time_step);
287:    if (th->endpoint) {
288:      TSComputeIJacobian(ts,ts->ptime,ts->vec_sol,th->Xdot,shift,J,Jp,PETSC_FALSE);
289:    }else {
290:      TSComputeIJacobian(ts,th->stage_time,th->X,th->Xdot,shift,J,Jp,PETSC_FALSE);
291:    }
292:    KSPSetOperators(ksp,J,Jp);

294:    /* Solve LHS X = RHS */
295:    for (nadj=0; nadj<ts->numcost; nadj++) {
296:      KSPSolveTranspose(ksp,VecsSensiTemp[nadj],VecsDeltaLam[nadj]);
297:    }

299:    /* Update sensitivities, and evaluate integrals if there is any */
300:    if(th->endpoint) { /* two-stage case */
301:      if (th->Theta!=1.) {
302:        shift = -1./((th->Theta-1.)*ts->time_step);
303:        TSComputeIJacobian(ts,th->ptime,th->X0,th->Xdot,shift,J,Jp,PETSC_FALSE);
304:        if (ts->vec_costintegral) {
305:          TSAdjointComputeDRDYFunction(ts,th->ptime,th->X0,ts->vecs_drdy);
306:        }
307:        for (nadj=0; nadj<ts->numcost; nadj++) {
308:          MatMultTranspose(J,VecsDeltaLam[nadj],ts->vecs_sensi[nadj]);
309:          if (ts->vec_costintegral) {
310:            VecAXPY(ts->vecs_sensi[nadj],-1.,ts->vecs_drdy[nadj]);
311:          }
312:          VecScale(ts->vecs_sensi[nadj],1./shift);
313:        }
314:      }else { /* backward Euler */
315:        shift = 0.0;
316:        TSComputeIJacobian(ts,ts->ptime,ts->vec_sol,th->Xdot,shift,J,Jp,PETSC_FALSE); /* get -f_y */
317:        for (nadj=0; nadj<ts->numcost; nadj++) {
318:          MatMultTranspose(J,VecsDeltaLam[nadj],VecsSensiTemp[nadj]);
319:          VecAXPY(ts->vecs_sensi[nadj],ts->time_step,VecsSensiTemp[nadj]);
320:          if (ts->vec_costintegral) {
321:            VecAXPY(ts->vecs_sensi[nadj],-ts->time_step,ts->vecs_drdy[nadj]);
322:          }
323:        }
324:      }

326:      if (ts->vecs_sensip) { /* sensitivities wrt parameters */
327:        TSAdjointComputeRHSJacobian(ts,ts->ptime,ts->vec_sol,ts->Jacp);
328:        for (nadj=0; nadj<ts->numcost; nadj++) {
329:          MatMultTranspose(ts->Jacp,VecsDeltaLam[nadj],VecsDeltaMu[nadj]);
330:          VecAXPY(ts->vecs_sensip[nadj],-ts->time_step*th->Theta,VecsDeltaMu[nadj]);
331:        }
332:        if (th->Theta!=1.) {
333:          TSAdjointComputeRHSJacobian(ts,th->ptime,th->X0,ts->Jacp);
334:          for (nadj=0; nadj<ts->numcost; nadj++) {
335:            MatMultTranspose(ts->Jacp,VecsDeltaLam[nadj],VecsDeltaMu[nadj]);
336:            VecAXPY(ts->vecs_sensip[nadj],-ts->time_step*(1.-th->Theta),VecsDeltaMu[nadj]);
337:          }
338:        }
339:        if (ts->vec_costintegral) {
340:          TSAdjointComputeDRDPFunction(ts,ts->ptime,ts->vec_sol,ts->vecs_drdp);
341:          for (nadj=0; nadj<ts->numcost; nadj++) {
342:            VecAXPY(ts->vecs_sensip[nadj],-ts->time_step*th->Theta,ts->vecs_drdp[nadj]);
343:          }
344:          if (th->Theta!=1.) {
345:            TSAdjointComputeDRDPFunction(ts,th->ptime,th->X0,ts->vecs_drdp);
346:            for (nadj=0; nadj<ts->numcost; nadj++) {
347:              VecAXPY(ts->vecs_sensip[nadj],-ts->time_step*(1.-th->Theta),ts->vecs_drdp[nadj]);
348:            }
349:          }
350:        }
351:      }
352:    }else { /* one-stage case */
353:      shift = 0.0;
354:      TSComputeIJacobian(ts,th->stage_time,th->X,th->Xdot,shift,J,Jp,PETSC_FALSE); /* get -f_y */
355:      if (ts->vec_costintegral) {
356:        TSAdjointComputeDRDYFunction(ts,th->stage_time,th->X,ts->vecs_drdy);
357:      }
358:      for (nadj=0; nadj<ts->numcost; nadj++) {
359:        MatMultTranspose(J,VecsDeltaLam[nadj],VecsSensiTemp[nadj]);
360:        VecAXPY(ts->vecs_sensi[nadj],ts->time_step,VecsSensiTemp[nadj]);
361:        if (ts->vec_costintegral) {
362:          VecAXPY(ts->vecs_sensi[nadj],-ts->time_step,ts->vecs_drdy[nadj]);
363:        }
364:      }
365:      if (ts->vecs_sensip) {
366:        TSAdjointComputeRHSJacobian(ts,th->stage_time,th->X,ts->Jacp);
367:        for (nadj=0; nadj<ts->numcost; nadj++) {
368:          MatMultTranspose(ts->Jacp,VecsDeltaLam[nadj],VecsDeltaMu[nadj]);
369:          VecAXPY(ts->vecs_sensip[nadj],-ts->time_step,VecsDeltaMu[nadj]);
370:        }
371:        if (ts->vec_costintegral) {
372:          TSAdjointComputeDRDPFunction(ts,th->stage_time,th->X,ts->vecs_drdp);
373:          for (nadj=0; nadj<ts->numcost; nadj++) {
374:            VecAXPY(ts->vecs_sensip[nadj],-ts->time_step,ts->vecs_drdp[nadj]);
375:          }
376:        }
377:      }
378:    }

380:    th->status = TS_STEP_COMPLETE;
381:    return(0);
382:  }

384:  static PetscErrorCode TSInterpolate_Theta(TS ts,PetscReal t,Vec X)
385:  {
386:    TS_Theta       *th = (TS_Theta*)ts->data;
387:    PetscReal      dt  = t - ts->ptime;

391:    VecCopy(ts->vec_sol,th->X);
392:    if (th->endpoint) dt *= th->Theta;
393:    VecWAXPY(X,dt,th->Xdot,th->X);
394:    return(0);
395:  }

397: static PetscErrorCode TSEvaluateWLTE_Theta(TS ts,NormType wnormtype,PetscInt *order,PetscReal *wlte)
398: {
399:   TS_Theta       *th = (TS_Theta*)ts->data;
400:   Vec            X = ts->vec_sol;      /* X = solution */
401:   Vec            Y = th->vec_lte_work; /* Y = X + LTE  */
402:   PetscReal      wltea,wlter;

406:   if (!th->vec_sol_prev) {*wlte = -1; return(0);}
407:   /* Cannot compute LTE in first step or in restart after event */
408:   if (ts->steprestart) {*wlte = -1; return(0);}
409:   /* Compute LTE using backward differences with non-constant time step */
410:   {
411:     PetscReal   h = ts->time_step, h_prev = ts->ptime - ts->ptime_prev;
412:     PetscReal   a = 1 + h_prev/h;
413:     PetscScalar scal[3]; Vec vecs[3];
414:     scal[0] = +1/a; scal[1] = -1/(a-1); scal[2] = +1/(a*(a-1));
415:     vecs[0] = X;    vecs[1] = th->X0;   vecs[2] = th->vec_sol_prev;
416:     VecCopy(X,Y);
417:     VecMAXPY(Y,3,scal,vecs);
418:     TSErrorWeightedNorm(ts,X,Y,wnormtype,wlte,&wltea,&wlter);
419:   }
420:   if (order) *order = 2;
421:   return(0);
422: }

424: static PetscErrorCode TSRollBack_Theta(TS ts)
425: {
426:   TS_Theta       *th = (TS_Theta*)ts->data;

430:   VecCopy(th->X0,ts->vec_sol);
431:   if (ts->vec_costintegral && ts->costintegralfwd) {
432:     VecCopy(th->VecCostIntegral0,ts->vec_costintegral);
433:   }
434:   return(0);
435: }

437: /*------------------------------------------------------------*/
438: static PetscErrorCode TSReset_Theta(TS ts)
439: {
440:   TS_Theta       *th = (TS_Theta*)ts->data;

444:   VecDestroy(&th->X);
445:   VecDestroy(&th->Xdot);
446:   VecDestroy(&th->X0);
447:   VecDestroy(&th->affine);

449:   VecDestroy(&th->vec_sol_prev);
450:   VecDestroy(&th->vec_lte_work);

452:   VecDestroy(&th->VecCostIntegral0);
453:   VecDestroyVecs(ts->numcost,&th->VecsDeltaLam);
454:   VecDestroyVecs(ts->numcost,&th->VecsDeltaMu);
455:   VecDestroyVecs(ts->numcost,&th->VecsSensiTemp);
456:   return(0);
457: }

459: static PetscErrorCode TSDestroy_Theta(TS ts)
460: {

464:   TSReset_Theta(ts);
465:   PetscFree(ts->data);
466:   PetscObjectComposeFunction((PetscObject)ts,"TSThetaGetTheta_C",NULL);
467:   PetscObjectComposeFunction((PetscObject)ts,"TSThetaSetTheta_C",NULL);
468:   PetscObjectComposeFunction((PetscObject)ts,"TSThetaGetEndpoint_C",NULL);
469:   PetscObjectComposeFunction((PetscObject)ts,"TSThetaSetEndpoint_C",NULL);
470:   return(0);
471: }

473: /*
474:   This defines the nonlinear equation that is to be solved with SNES
475:   G(U) = F[t0+Theta*dt, U, (U-U0)*shift] = 0
476: */
477: static PetscErrorCode SNESTSFormFunction_Theta(SNES snes,Vec x,Vec y,TS ts)
478: {
479:   TS_Theta       *th = (TS_Theta*)ts->data;
481:   Vec            X0,Xdot;
482:   DM             dm,dmsave;
483:   PetscReal      shift = 1/(th->Theta*ts->time_step);

486:   SNESGetDM(snes,&dm);
487:   /* When using the endpoint variant, this is actually 1/Theta * Xdot */
488:   TSThetaGetX0AndXdot(ts,dm,&X0,&Xdot);
489:   VecAXPBYPCZ(Xdot,-shift,shift,0,X0,x);

491:   /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */
492:   dmsave = ts->dm;
493:   ts->dm = dm;
494:   TSComputeIFunction(ts,th->stage_time,x,Xdot,y,PETSC_FALSE);
495:   ts->dm = dmsave;
496:   TSThetaRestoreX0AndXdot(ts,dm,&X0,&Xdot);
497:   return(0);
498: }

500: static PetscErrorCode SNESTSFormJacobian_Theta(SNES snes,Vec x,Mat A,Mat B,TS ts)
501: {
502:   TS_Theta       *th = (TS_Theta*)ts->data;
504:   Vec            Xdot;
505:   DM             dm,dmsave;
506:   PetscReal      shift = 1/(th->Theta*ts->time_step);

509:   SNESGetDM(snes,&dm);
510:   /* Xdot has already been computed in SNESTSFormFunction_Theta (SNES guarantees this) */
511:   TSThetaGetX0AndXdot(ts,dm,NULL,&Xdot);

513:   dmsave = ts->dm;
514:   ts->dm = dm;
515:   TSComputeIJacobian(ts,th->stage_time,x,Xdot,shift,A,B,PETSC_FALSE);
516:   ts->dm = dmsave;
517:   TSThetaRestoreX0AndXdot(ts,dm,NULL,&Xdot);
518:   return(0);
519: }

521: static PetscErrorCode TSSetUp_Theta(TS ts)
522: {
523:   TS_Theta       *th = (TS_Theta*)ts->data;
524:   PetscBool      match;

528:   if (!th->VecCostIntegral0 && ts->vec_costintegral && ts->costintegralfwd) { /* back up cost integral */
529:     VecDuplicate(ts->vec_costintegral,&th->VecCostIntegral0);
530:   }
531:   if (!th->X) {
532:     VecDuplicate(ts->vec_sol,&th->X);
533:   }
534:   if (!th->Xdot) {
535:     VecDuplicate(ts->vec_sol,&th->Xdot);
536:   }
537:   if (!th->X0) {
538:     VecDuplicate(ts->vec_sol,&th->X0);
539:   }
540:   if (th->endpoint) {
541:     VecDuplicate(ts->vec_sol,&th->affine);
542:   }

544:   th->order = (th->Theta == 0.5) ? 2 : 1;

546:   TSGetDM(ts,&ts->dm);
547:   DMCoarsenHookAdd(ts->dm,DMCoarsenHook_TSTheta,DMRestrictHook_TSTheta,ts);
548:   DMSubDomainHookAdd(ts->dm,DMSubDomainHook_TSTheta,DMSubDomainRestrictHook_TSTheta,ts);

550:   TSGetAdapt(ts,&ts->adapt);
551:   TSAdaptCandidatesClear(ts->adapt);
552:   PetscObjectTypeCompare((PetscObject)ts->adapt,TSADAPTNONE,&match);
553:   if (!match) {
554:     VecDuplicate(ts->vec_sol,&th->vec_sol_prev);
555:     VecDuplicate(ts->vec_sol,&th->vec_lte_work);
556:   }

558:   TSGetSNES(ts,&ts->snes);
559:   return(0);
560: }

562: /*------------------------------------------------------------*/

564: static PetscErrorCode TSAdjointSetUp_Theta(TS ts)
565: {
566:   TS_Theta       *th = (TS_Theta*)ts->data;

570:   VecDuplicateVecs(ts->vecs_sensi[0],ts->numcost,&th->VecsDeltaLam);
571:   if(ts->vecs_sensip) {
572:     VecDuplicateVecs(ts->vecs_sensip[0],ts->numcost,&th->VecsDeltaMu);
573:   }
574:   VecDuplicateVecs(ts->vecs_sensi[0],ts->numcost,&th->VecsSensiTemp);
575:   return(0);
576: }
577: /*------------------------------------------------------------*/

579: static PetscErrorCode TSSetFromOptions_Theta(PetscOptionItems *PetscOptionsObject,TS ts)
580: {
581:   TS_Theta       *th = (TS_Theta*)ts->data;

585:   PetscOptionsHead(PetscOptionsObject,"Theta ODE solver options");
586:   {
587:     PetscOptionsReal("-ts_theta_theta","Location of stage (0<Theta<=1)","TSThetaSetTheta",th->Theta,&th->Theta,NULL);
588:     PetscOptionsBool("-ts_theta_endpoint","Use the endpoint instead of midpoint form of the Theta method","TSThetaSetEndpoint",th->endpoint,&th->endpoint,NULL);
589:     PetscOptionsBool("-ts_theta_initial_guess_extrapolate","Extrapolate stage initial guess from previous solution (sometimes unstable)","TSThetaSetExtrapolate",th->extrapolate,&th->extrapolate,NULL);
590:   }
591:   PetscOptionsTail();
592:   return(0);
593: }

595: static PetscErrorCode TSView_Theta(TS ts,PetscViewer viewer)
596: {
597:   TS_Theta       *th = (TS_Theta*)ts->data;
598:   PetscBool      iascii;

602:   PetscObjectTypeCompare((PetscObject)viewer,PETSCVIEWERASCII,&iascii);
603:   if (iascii) {
604:     PetscViewerASCIIPrintf(viewer,"  Theta=%g\n",(double)th->Theta);
605:     PetscViewerASCIIPrintf(viewer,"  Extrapolation=%s\n",th->extrapolate ? "yes" : "no");
606:   }
607:   return(0);
608: }

610: static PetscErrorCode TSThetaGetTheta_Theta(TS ts,PetscReal *theta)
611: {
612:   TS_Theta *th = (TS_Theta*)ts->data;

615:   *theta = th->Theta;
616:   return(0);
617: }

619: static PetscErrorCode TSThetaSetTheta_Theta(TS ts,PetscReal theta)
620: {
621:   TS_Theta *th = (TS_Theta*)ts->data;

624:   if (theta <= 0 || 1 < theta) SETERRQ1(PetscObjectComm((PetscObject)ts),PETSC_ERR_ARG_OUTOFRANGE,"Theta %g not in range (0,1]",(double)theta);
625:   th->Theta = theta;
626:   th->order = (th->Theta == 0.5) ? 2 : 1;
627:   return(0);
628: }

630: static PetscErrorCode TSThetaGetEndpoint_Theta(TS ts,PetscBool *endpoint)
631: {
632:   TS_Theta *th = (TS_Theta*)ts->data;

635:   *endpoint = th->endpoint;
636:   return(0);
637: }

639: static PetscErrorCode TSThetaSetEndpoint_Theta(TS ts,PetscBool flg)
640: {
641:   TS_Theta *th = (TS_Theta*)ts->data;

644:   th->endpoint = flg;
645:   return(0);
646: }

648: #if defined(PETSC_HAVE_COMPLEX)
649: static PetscErrorCode TSComputeLinearStability_Theta(TS ts,PetscReal xr,PetscReal xi,PetscReal *yr,PetscReal *yi)
650: {
651:   PetscComplex   z   = xr + xi*PETSC_i,f;
652:   TS_Theta       *th = (TS_Theta*)ts->data;
653:   const PetscReal one = 1.0;

656:   f   = (one + (one - th->Theta)*z)/(one - th->Theta*z);
657:   *yr = PetscRealPartComplex(f);
658:   *yi = PetscImaginaryPartComplex(f);
659:   return(0);
660: }
661: #endif

663: static PetscErrorCode TSGetStages_Theta(TS ts,PetscInt *ns,Vec **Y)
664: {
665:   TS_Theta     *th = (TS_Theta*)ts->data;

668:   if (ns) *ns = 1;
669:   if (Y)  *Y  = th->endpoint ? &(th->X0) : &(th->X);
670:   return(0);
671: }

673: /* ------------------------------------------------------------ */
674: /*MC
675:       TSTHETA - DAE solver using the implicit Theta method

677:    Level: beginner

679:    Options Database:
680: +  -ts_theta_theta <Theta> - Location of stage (0<Theta<=1)
681: .  -ts_theta_endpoint <flag> - Use the endpoint (like Crank-Nicholson) instead of midpoint form of the Theta method
682: -  -ts_theta_initial_guess_extrapolate <flg> - Extrapolate stage initial guess from previous solution (sometimes unstable)

684:    Notes:
685: $  -ts_type theta -ts_theta_theta 1.0 corresponds to backward Euler (TSBEULER)
686: $  -ts_type theta -ts_theta_theta 0.5 corresponds to the implicit midpoint rule
687: $  -ts_type theta -ts_theta_theta 0.5 -ts_theta_endpoint corresponds to Crank-Nicholson (TSCN)

689:    This method can be applied to DAE.

691:    This method is cast as a 1-stage implicit Runge-Kutta method.

693: .vb
694:   Theta | Theta
695:   -------------
696:         |  1
697: .ve

699:    For the default Theta=0.5, this is also known as the implicit midpoint rule.

701:    When the endpoint variant is chosen, the method becomes a 2-stage method with first stage explicit:

703: .vb
704:   0 | 0         0
705:   1 | 1-Theta   Theta
706:   -------------------
707:     | 1-Theta   Theta
708: .ve

710:    For the default Theta=0.5, this is the trapezoid rule (also known as Crank-Nicolson, see TSCN).

712:    To apply a diagonally implicit RK method to DAE, the stage formula

714: $  Y_i = X + h sum_j a_ij Y'_j

716:    is interpreted as a formula for Y'_i in terms of Y_i and known values (Y'_j, j<i)

718: .seealso:  TSCreate(), TS, TSSetType(), TSCN, TSBEULER, TSThetaSetTheta(), TSThetaSetEndpoint()

720: M*/
721: PETSC_EXTERN PetscErrorCode TSCreate_Theta(TS ts)
722: {
723:   TS_Theta       *th;

727:   ts->ops->reset           = TSReset_Theta;
728:   ts->ops->destroy         = TSDestroy_Theta;
729:   ts->ops->view            = TSView_Theta;
730:   ts->ops->setup           = TSSetUp_Theta;
731:   ts->ops->adjointsetup    = TSAdjointSetUp_Theta;
732:   ts->ops->step            = TSStep_Theta;
733:   ts->ops->interpolate     = TSInterpolate_Theta;
734:   ts->ops->evaluatewlte    = TSEvaluateWLTE_Theta;
735:   ts->ops->rollback        = TSRollBack_Theta;
736:   ts->ops->setfromoptions  = TSSetFromOptions_Theta;
737:   ts->ops->snesfunction    = SNESTSFormFunction_Theta;
738:   ts->ops->snesjacobian    = SNESTSFormJacobian_Theta;
739: #if defined(PETSC_HAVE_COMPLEX)
740:   ts->ops->linearstability = TSComputeLinearStability_Theta;
741: #endif
742:   ts->ops->getstages       = TSGetStages_Theta;
743:   ts->ops->adjointstep     = TSAdjointStep_Theta;
744:   ts->ops->adjointintegral = TSAdjointCostIntegral_Theta;
745:   ts->ops->forwardintegral = TSForwardCostIntegral_Theta;
746:   ts->default_adapt_type   = TSADAPTNONE;

748:   ts->usessnes = PETSC_TRUE;

750:   PetscNewLog(ts,&th);
751:   ts->data = (void*)th;

753:   th->extrapolate = PETSC_FALSE;
754:   th->Theta       = 0.5;
755:   th->order       = 2;
756:   PetscObjectComposeFunction((PetscObject)ts,"TSThetaGetTheta_C",TSThetaGetTheta_Theta);
757:   PetscObjectComposeFunction((PetscObject)ts,"TSThetaSetTheta_C",TSThetaSetTheta_Theta);
758:   PetscObjectComposeFunction((PetscObject)ts,"TSThetaGetEndpoint_C",TSThetaGetEndpoint_Theta);
759:   PetscObjectComposeFunction((PetscObject)ts,"TSThetaSetEndpoint_C",TSThetaSetEndpoint_Theta);
760:   return(0);
761: }

763: /*@
764:   TSThetaGetTheta - Get the abscissa of the stage in (0,1].

766:   Not Collective

768:   Input Parameter:
769: .  ts - timestepping context

771:   Output Parameter:
772: .  theta - stage abscissa

774:   Note:
775:   Use of this function is normally only required to hack TSTHETA to use a modified integration scheme.

777:   Level: Advanced

779: .seealso: TSThetaSetTheta()
780: @*/
781: PetscErrorCode  TSThetaGetTheta(TS ts,PetscReal *theta)
782: {

788:   PetscUseMethod(ts,"TSThetaGetTheta_C",(TS,PetscReal*),(ts,theta));
789:   return(0);
790: }

792: /*@
793:   TSThetaSetTheta - Set the abscissa of the stage in (0,1].

795:   Not Collective

797:   Input Parameter:
798: +  ts - timestepping context
799: -  theta - stage abscissa

801:   Options Database:
802: .  -ts_theta_theta <theta>

804:   Level: Intermediate

806: .seealso: TSThetaGetTheta()
807: @*/
808: PetscErrorCode  TSThetaSetTheta(TS ts,PetscReal theta)
809: {

814:   PetscTryMethod(ts,"TSThetaSetTheta_C",(TS,PetscReal),(ts,theta));
815:   return(0);
816: }

818: /*@
819:   TSThetaGetEndpoint - Gets whether to use the endpoint variant of the method (e.g. trapezoid/Crank-Nicolson instead of midpoint rule).

821:   Not Collective

823:   Input Parameter:
824: .  ts - timestepping context

826:   Output Parameter:
827: .  endpoint - PETSC_TRUE when using the endpoint variant

829:   Level: Advanced

831: .seealso: TSThetaSetEndpoint(), TSTHETA, TSCN
832: @*/
833: PetscErrorCode TSThetaGetEndpoint(TS ts,PetscBool *endpoint)
834: {

840:   PetscUseMethod(ts,"TSThetaGetEndpoint_C",(TS,PetscBool*),(ts,endpoint));
841:   return(0);
842: }

844: /*@
845:   TSThetaSetEndpoint - Sets whether to use the endpoint variant of the method (e.g. trapezoid/Crank-Nicolson instead of midpoint rule).

847:   Not Collective

849:   Input Parameter:
850: +  ts - timestepping context
851: -  flg - PETSC_TRUE to use the endpoint variant

853:   Options Database:
854: .  -ts_theta_endpoint <flg>

856:   Level: Intermediate

858: .seealso: TSTHETA, TSCN
859: @*/
860: PetscErrorCode TSThetaSetEndpoint(TS ts,PetscBool flg)
861: {

866:   PetscTryMethod(ts,"TSThetaSetEndpoint_C",(TS,PetscBool),(ts,flg));
867:   return(0);
868: }

870: /*
871:  * TSBEULER and TSCN are straightforward specializations of TSTHETA.
872:  * The creation functions for these specializations are below.
873:  */

875: static PetscErrorCode TSSetUp_BEuler(TS ts)
876: {
877:   TS_Theta       *th = (TS_Theta*)ts->data;

881:   if (th->Theta != 1.0) SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_OPT_OVERWRITE,"Can not change the default value (1) of theta when using backward Euler\n");
882:   if (th->endpoint) SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_OPT_OVERWRITE,"Can not change to the endpoint form of the Theta methods when using backward Euler\n");
883:   TSSetUp_Theta(ts);
884:   return(0);
885: }

887: static PetscErrorCode TSView_BEuler(TS ts,PetscViewer viewer)
888: {
890:   return(0);
891: }

893: /*MC
894:       TSBEULER - ODE solver using the implicit backward Euler method

896:   Level: beginner

898:   Notes:
899:   TSBEULER is equivalent to TSTHETA with Theta=1.0

901: $  -ts_type theta -ts_theta_theta 1.0

903: .seealso:  TSCreate(), TS, TSSetType(), TSEULER, TSCN, TSTHETA

905: M*/
906: PETSC_EXTERN PetscErrorCode TSCreate_BEuler(TS ts)
907: {

911:   TSCreate_Theta(ts);
912:   TSThetaSetTheta(ts,1.0);
913:   TSThetaSetEndpoint(ts,PETSC_FALSE);
914:   ts->ops->setup = TSSetUp_BEuler;
915:   ts->ops->view  = TSView_BEuler;
916:   return(0);
917: }

919: static PetscErrorCode TSSetUp_CN(TS ts)
920: {
921:   TS_Theta       *th = (TS_Theta*)ts->data;

925:   if (th->Theta != 0.5) SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_OPT_OVERWRITE,"Can not change the default value (0.5) of theta when using Crank-Nicolson\n");
926:   if (!th->endpoint) SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_OPT_OVERWRITE,"Can not change to the midpoint form of the Theta methods when using Crank-Nicolson\n");
927:   TSSetUp_Theta(ts);
928:   return(0);
929: }

931: static PetscErrorCode TSView_CN(TS ts,PetscViewer viewer)
932: {
934:   return(0);
935: }

937: /*MC
938:       TSCN - ODE solver using the implicit Crank-Nicolson method.

940:   Level: beginner

942:   Notes:
943:   TSCN is equivalent to TSTHETA with Theta=0.5 and the "endpoint" option set. I.e.

945: $  -ts_type theta -ts_theta_theta 0.5 -ts_theta_endpoint

947: .seealso:  TSCreate(), TS, TSSetType(), TSBEULER, TSTHETA

949: M*/
950: PETSC_EXTERN PetscErrorCode TSCreate_CN(TS ts)
951: {

955:   TSCreate_Theta(ts);
956:   TSThetaSetTheta(ts,0.5);
957:   TSThetaSetEndpoint(ts,PETSC_TRUE);
958:   ts->ops->setup = TSSetUp_CN;
959:   ts->ops->view  = TSView_CN;
960:   return(0);
961: }