Date: 2015-10-02T22:27:03+00:00

Actual source code: theta.c
petsc-3.6.2 2015-10-02
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:   Vec          X,Xdot;                   /* Storage for one stage */
     11:   Vec          X0;                       /* work vector to store X0 */
     12:   Vec          affine;                   /* Affine vector needed for residual at beginning of step */
     13:   Vec          *VecsDeltaLam;             /* Increment of the adjoint sensitivity w.r.t IC at stage*/
     14:   Vec          *VecsDeltaMu;              /* Increment of the adjoint sensitivity w.r.t P at stage*/
     15:   Vec          *VecsSensiTemp;            /* Vector to be timed with Jacobian transpose*/
     16:   PetscBool    extrapolate;
     17:   PetscBool    endpoint;
     18:   PetscReal    Theta;
     19:   PetscReal    stage_time;
     20:   TSStepStatus status;
     21:   char         *name;
     22:   PetscInt     order;
     23:   PetscReal    ccfl;               /* Placeholder for CFL coefficient relative to forward Euler */
     24:   PetscBool    adapt;  /* use time-step adaptivity ? */
     25:   PetscReal    ptime;
     26: } TS_Theta;

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

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

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

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

     72: static PetscErrorCode DMCoarsenHook_TSTheta(DM fine,DM coarse,void *ctx)
     73: {

     76:   return(0);
     77: }

     81: static PetscErrorCode DMRestrictHook_TSTheta(DM fine,Mat restrct,Vec rscale,Mat inject,DM coarse,void *ctx)
     82: {
     83:   TS             ts = (TS)ctx;
     85:   Vec            X0,Xdot,X0_c,Xdot_c;

     88:   TSThetaGetX0AndXdot(ts,fine,&X0,&Xdot);
     89:   TSThetaGetX0AndXdot(ts,coarse,&X0_c,&Xdot_c);
     90:   MatRestrict(restrct,X0,X0_c);
     91:   MatRestrict(restrct,Xdot,Xdot_c);
     92:   VecPointwiseMult(X0_c,rscale,X0_c);
     93:   VecPointwiseMult(Xdot_c,rscale,Xdot_c);
     94:   TSThetaRestoreX0AndXdot(ts,fine,&X0,&Xdot);
     95:   TSThetaRestoreX0AndXdot(ts,coarse,&X0_c,&Xdot_c);
     96:   return(0);
     97: }

    101: static PetscErrorCode DMSubDomainHook_TSTheta(DM dm,DM subdm,void *ctx)
    102: {

    105:   return(0);
    106: }

    110: static PetscErrorCode DMSubDomainRestrictHook_TSTheta(DM dm,VecScatter gscat,VecScatter lscat,DM subdm,void *ctx)
    111: {
    112:   TS             ts = (TS)ctx;
    114:   Vec            X0,Xdot,X0_sub,Xdot_sub;

    117:   TSThetaGetX0AndXdot(ts,dm,&X0,&Xdot);
    118:   TSThetaGetX0AndXdot(ts,subdm,&X0_sub,&Xdot_sub);

    120:   VecScatterBegin(gscat,X0,X0_sub,INSERT_VALUES,SCATTER_FORWARD);
    121:   VecScatterEnd(gscat,X0,X0_sub,INSERT_VALUES,SCATTER_FORWARD);

    123:   VecScatterBegin(gscat,Xdot,Xdot_sub,INSERT_VALUES,SCATTER_FORWARD);
    124:   VecScatterEnd(gscat,Xdot,Xdot_sub,INSERT_VALUES,SCATTER_FORWARD);

    126:   TSThetaRestoreX0AndXdot(ts,dm,&X0,&Xdot);
    127:   TSThetaRestoreX0AndXdot(ts,subdm,&X0_sub,&Xdot_sub);
    128:   return(0);
    129: }

    133: static PetscErrorCode TSEvaluateStep_Theta(TS ts,PetscInt order,Vec U,PetscBool *done)
    134: {
    136:   TS_Theta       *th = (TS_Theta*)ts->data;

    139:   if (order == 0) SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_USER,"No time-step adaptivity implemented for 1st order theta method; Run with -ts_adapt_type none");
    140:   if (order == th->order) {
    141:     if (th->endpoint) {
    142:       VecCopy(th->X,U);
    143:     } else {
    144:       PetscReal shift = 1./(th->Theta*ts->time_step);
    145:       VecAXPBYPCZ(th->Xdot,-shift,shift,0,U,th->X);
    146:       VecAXPY(U,ts->time_step,th->Xdot);
    147:     }
    148:   } else if (order == th->order-1 && order) {
    149:     VecWAXPY(U,ts->time_step,th->Xdot,th->X0);
    150:   }
    151:   return(0);
    152: }

    156: static PetscErrorCode TSRollBack_Theta(TS ts)
    157: {
    158:   TS_Theta       *th = (TS_Theta*)ts->data;

    162:   VecCopy(th->X0,ts->vec_sol);
    163:   th->status    = TS_STEP_INCOMPLETE;
    164:   return(0);
    165: }

    169: static PetscErrorCode TSStep_Theta(TS ts)
    170: {
    171:   TS_Theta       *th = (TS_Theta*)ts->data;
    172:   PetscInt       its,lits,reject,next_scheme;
    173:   PetscReal      next_time_step;
    174:   TSAdapt        adapt;
    175:   PetscBool      stageok,accept = PETSC_TRUE;

    179:   th->status = TS_STEP_INCOMPLETE;
    180:   VecCopy(ts->vec_sol,th->X0);
    181:   for (reject=0; !ts->reason && th->status != TS_STEP_COMPLETE; ts->reject++) {
    182:     PetscReal shift = 1./(th->Theta*ts->time_step);
    183:     th->stage_time = ts->ptime + (th->endpoint ? 1. : th->Theta)*ts->time_step;
    184:     TSPreStep(ts);
    185:     TSPreStage(ts,th->stage_time);

    187:     if (th->endpoint) {           /* This formulation assumes linear time-independent mass matrix */
    188:       VecZeroEntries(th->Xdot);
    189:       if (!th->affine) {VecDuplicate(ts->vec_sol,&th->affine);}
    190:       TSComputeIFunction(ts,ts->ptime,ts->vec_sol,th->Xdot,th->affine,PETSC_FALSE);
    191:       VecScale(th->affine,(th->Theta-1.)/th->Theta);
    192:     }
    193:     if (th->extrapolate) {
    194:       VecWAXPY(th->X,1./shift,th->Xdot,ts->vec_sol);
    195:     } else {
    196:       VecCopy(ts->vec_sol,th->X);
    197:     }
    198:     SNESSolve(ts->snes,th->affine,th->X);
    199:     SNESGetIterationNumber(ts->snes,&its);
    200:     SNESGetLinearSolveIterations(ts->snes,&lits);
    201:     ts->snes_its += its; ts->ksp_its += lits;
    202:     TSPostStage(ts,th->stage_time,0,&(th->X));
    203:     TSGetAdapt(ts,&adapt);
    204:     TSAdaptCheckStage(adapt,ts,&stageok);
    205:     if (!stageok) {accept = PETSC_FALSE; goto reject_step;}

    207:     TSEvaluateStep(ts,th->order,ts->vec_sol,NULL);
    208:     th->status = TS_STEP_PENDING;
    209:     /* Register only the current method as a candidate because we're not supporting multiple candidates yet. */
    210:     TSGetAdapt(ts,&adapt);
    211:     TSAdaptCandidatesClear(adapt);
    212:     TSAdaptCandidateAdd(adapt,NULL,th->order,1,th->ccfl,1.0,PETSC_TRUE);
    213:     TSAdaptChoose(adapt,ts,ts->time_step,&next_scheme,&next_time_step,&accept);
    214:     if (!accept) {           /* Roll back the current step */
    215:       ts->ptime += next_time_step; /* This will be undone in rollback */
    216:       th->status = TS_STEP_INCOMPLETE;
    217:       TSRollBack(ts);
    218:       goto reject_step;
    219:     }

    221:     if (ts->vec_costintegral) {
    222:       /* Evolve ts->vec_costintegral to compute integrals */
    223:       if (th->endpoint) {
    224:         TSAdjointComputeCostIntegrand(ts,ts->ptime,th->X0,ts->vec_costintegrand);
    225:         VecAXPY(ts->vec_costintegral,ts->time_step*(1.-th->Theta),ts->vec_costintegrand);
    226:       }
    227:       TSAdjointComputeCostIntegrand(ts,th->stage_time,th->X,ts->vec_costintegrand);
    228:       if (th->endpoint) {
    229:         VecAXPY(ts->vec_costintegral,ts->time_step*th->Theta,ts->vec_costintegrand);
    230:       }else {
    231:         VecAXPY(ts->vec_costintegral,ts->time_step,ts->vec_costintegrand);
    232:       }
    233:     }

    235:     /* ignore next_scheme for now */
    236:     ts->ptime    += ts->time_step;
    237:     ts->time_step = next_time_step;
    238:     ts->steps++;
    239:     th->status = TS_STEP_COMPLETE;
    240:     break;

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

    254: static PetscErrorCode TSAdjointStep_Theta(TS ts)
    255: {
    256:   TS_Theta            *th = (TS_Theta*)ts->data;
    257:   Vec                 *VecsDeltaLam = th->VecsDeltaLam,*VecsDeltaMu = th->VecsDeltaMu,*VecsSensiTemp = th->VecsSensiTemp;
    258:   PetscInt            nadj;
    259:   PetscErrorCode      ierr;
    260:   Mat                 J,Jp;
    261:   KSP                 ksp;
    262:   PetscReal           shift;

    265:
    266:   th->status = TS_STEP_INCOMPLETE;
    267:   SNESGetKSP(ts->snes,&ksp);
    268:   TSGetIJacobian(ts,&J,&Jp,NULL,NULL);

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

    274:   TSPreStep(ts);

    276:   /* Build RHS */
    277:   if (ts->vec_costintegral) { /* Cost function has an integral term */
    278:     if (th->endpoint) {
    279:       TSAdjointComputeDRDYFunction(ts,ts->ptime,ts->vec_sol,ts->vecs_drdy);
    280:     }else {
    281:       TSAdjointComputeDRDYFunction(ts,th->stage_time,th->X,ts->vecs_drdy);
    282:     }
    283:   }
    284:   for (nadj=0; nadj<ts->numcost; nadj++) {
    285:     VecCopy(ts->vecs_sensi[nadj],VecsSensiTemp[nadj]);
    286:     VecScale(VecsSensiTemp[nadj],-1./(th->Theta*ts->time_step));
    287:     if (ts->vec_costintegral) {
    288:       VecAXPY(VecsSensiTemp[nadj],1.,ts->vecs_drdy[nadj]);
    289:     }
    290:   }
    291:
    292:   /* Build LHS */
    293:   shift = -1./(th->Theta*ts->time_step);
    294:   if (th->endpoint) {
    295:     TSComputeIJacobian(ts,ts->ptime,ts->vec_sol,th->Xdot,shift,J,Jp,PETSC_FALSE);
    296:   }else {
    297:     TSComputeIJacobian(ts,th->stage_time,th->X,th->Xdot,shift,J,Jp,PETSC_FALSE);
    298:   }
    299:   KSPSetOperators(ksp,J,Jp);

    301:   /* Solve LHS X = RHS */
    302:   for (nadj=0; nadj<ts->numcost; nadj++) {
    303:     KSPSolveTranspose(ksp,VecsSensiTemp[nadj],VecsDeltaLam[nadj]);
    304:   }

    306:   /* Update sensitivities, and evaluate integrals if there is any */
    307:   if(th->endpoint) { /* two-stage case */
    308:     if (th->Theta!=1.) {
    309:       shift = -1./((th->Theta-1.)*ts->time_step);
    310:       TSComputeIJacobian(ts,th->ptime,th->X0,th->Xdot,shift,J,Jp,PETSC_FALSE);
    311:       if (ts->vec_costintegral) {
    312:         TSAdjointComputeDRDYFunction(ts,th->ptime,th->X0,ts->vecs_drdy);
    313:         if (!ts->costintegralfwd) {
    314:           /* Evolve ts->vec_costintegral to compute integrals */
    315:           TSAdjointComputeCostIntegrand(ts,ts->ptime,ts->vec_sol,ts->vec_costintegrand);
    316:           VecAXPY(ts->vec_costintegral,-ts->time_step*th->Theta,ts->vec_costintegrand);
    317:           TSAdjointComputeCostIntegrand(ts,th->ptime,th->X0,ts->vec_costintegrand);
    318:           VecAXPY(ts->vec_costintegral,ts->time_step*(th->Theta-1.),ts->vec_costintegrand);
    319:         }
    320:       }
    321:       for (nadj=0; nadj<ts->numcost; nadj++) {
    322:         MatMultTranspose(J,VecsDeltaLam[nadj],ts->vecs_sensi[nadj]);
    323:         if (ts->vec_costintegral) {
    324:           VecAXPY(ts->vecs_sensi[nadj],-1.,ts->vecs_drdy[nadj]);
    325:         }
    326:         VecScale(ts->vecs_sensi[nadj],1./shift);
    327:       }
    328:     }else { /* backward Euler */
    329:       shift = 0.0;
    330:       TSComputeIJacobian(ts,ts->ptime,ts->vec_sol,th->Xdot,shift,J,Jp,PETSC_FALSE); /* get -f_y */
    331:       for (nadj=0; nadj<ts->numcost; nadj++) {
    332:         MatMultTranspose(J,VecsDeltaLam[nadj],VecsSensiTemp[nadj]);
    333:         VecAXPY(ts->vecs_sensi[nadj],ts->time_step,VecsSensiTemp[nadj]);
    334:         if (ts->vec_costintegral) {
    335:           VecAXPY(ts->vecs_sensi[nadj],-ts->time_step,ts->vecs_drdy[nadj]);
    336:           if (!ts->costintegralfwd) {
    337:           /* Evolve ts->vec_costintegral to compute integrals */
    338:             TSAdjointComputeCostIntegrand(ts,ts->ptime,ts->vec_sol,ts->vec_costintegrand);
    339:             VecAXPY(ts->vec_costintegral,-ts->time_step*th->Theta,ts->vec_costintegrand);
    340:           }
    341:         }
    342:       }
    343:     }

    345:     if (ts->vecs_sensip) { /* sensitivities wrt parameters */
    346:       TSAdjointComputeRHSJacobian(ts,ts->ptime,ts->vec_sol,ts->Jacp);
    347:       for (nadj=0; nadj<ts->numcost; nadj++) {
    348:         MatMultTranspose(ts->Jacp,VecsDeltaLam[nadj],VecsDeltaMu[nadj]);
    349:         VecAXPY(ts->vecs_sensip[nadj],-ts->time_step*th->Theta,VecsDeltaMu[nadj]);
    350:       }
    351:       if (th->Theta!=1.) {
    352:         TSAdjointComputeRHSJacobian(ts,th->ptime,th->X0,ts->Jacp);
    353:         for (nadj=0; nadj<ts->numcost; nadj++) {
    354:           MatMultTranspose(ts->Jacp,VecsDeltaLam[nadj],VecsDeltaMu[nadj]);
    355:           VecAXPY(ts->vecs_sensip[nadj],-ts->time_step*(1.-th->Theta),VecsDeltaMu[nadj]);
    356:         }
    357:       }
    358:       if (ts->vec_costintegral) {
    359:         TSAdjointComputeDRDPFunction(ts,ts->ptime,ts->vec_sol,ts->vecs_drdp);
    360:         for (nadj=0; nadj<ts->numcost; nadj++) {
    361:           VecAXPY(ts->vecs_sensip[nadj],-ts->time_step*th->Theta,ts->vecs_drdp[nadj]);
    362:         }
    363:         if (th->Theta!=1.) {
    364:           TSAdjointComputeDRDPFunction(ts,th->ptime,th->X0,ts->vecs_drdp);
    365:           for (nadj=0; nadj<ts->numcost; nadj++) {
    366:             VecAXPY(ts->vecs_sensip[nadj],-ts->time_step*(1.-th->Theta),ts->vecs_drdp[nadj]);
    367:           }
    368:         }
    369:       }
    370:     }
    371:   }else { /* one-stage case */
    372:     shift = 0.0;
    373:     TSComputeIJacobian(ts,th->stage_time,th->X,th->Xdot,shift,J,Jp,PETSC_FALSE); /* get -f_y */
    374:     if (ts->vec_costintegral) {
    375:       TSAdjointComputeDRDYFunction(ts,th->stage_time,th->X,ts->vecs_drdy);
    376:       if (!ts->costintegralfwd) {
    377:         /* Evolve ts->vec_costintegral to compute integrals */
    378:         TSAdjointComputeCostIntegrand(ts,th->stage_time,th->X,ts->vec_costintegrand);
    379:         VecAXPY(ts->vec_costintegral,-ts->time_step,ts->vec_costintegrand);
    380:       }
    381:     }
    382:     for (nadj=0; nadj<ts->numcost; nadj++) {
    383:       MatMultTranspose(J,VecsDeltaLam[nadj],VecsSensiTemp[nadj]);
    384:       VecAXPY(ts->vecs_sensi[nadj],ts->time_step,VecsSensiTemp[nadj]);
    385:       if (ts->vec_costintegral) {
    386:         VecAXPY(ts->vecs_sensi[nadj],-ts->time_step,ts->vecs_drdy[nadj]);
    387:       }
    388:     }
    389:     if (ts->vecs_sensip) {
    390:       TSAdjointComputeRHSJacobian(ts,th->stage_time,th->X,ts->Jacp);
    391:       for (nadj=0; nadj<ts->numcost; nadj++) {
    392:         MatMultTranspose(ts->Jacp,VecsDeltaLam[nadj],VecsDeltaMu[nadj]);
    393:         VecAXPY(ts->vecs_sensip[nadj],-ts->time_step,VecsDeltaMu[nadj]);
    394:       }
    395:       if (ts->vec_costintegral) {
    396:         TSAdjointComputeDRDPFunction(ts,th->stage_time,th->X,ts->vecs_drdp);
    397:         for (nadj=0; nadj<ts->numcost; nadj++) {
    398:           VecAXPY(ts->vecs_sensip[nadj],-ts->time_step,ts->vecs_drdp[nadj]);
    399:         }
    400:       }
    401:     }
    402:   }
    403:
    404:   ts->ptime += ts->time_step;
    405:   ts->steps++;
    406:   th->status = TS_STEP_COMPLETE;
    407:   return(0);
    408: }

    412: static PetscErrorCode TSInterpolate_Theta(TS ts,PetscReal t,Vec X)
    413: {
    414:   TS_Theta       *th   = (TS_Theta*)ts->data;
    415:   PetscReal      alpha = t - ts->ptime;

    419:   VecCopy(ts->vec_sol,th->X);
    420:   if (th->endpoint) alpha *= th->Theta;
    421:   VecWAXPY(X,alpha,th->Xdot,th->X);
    422:   return(0);
    423: }

    425: /*------------------------------------------------------------*/
    428: static PetscErrorCode TSReset_Theta(TS ts)
    429: {
    430:   TS_Theta       *th = (TS_Theta*)ts->data;

    434:   VecDestroy(&th->X);
    435:   VecDestroy(&th->Xdot);
    436:   VecDestroy(&th->X0);
    437:   VecDestroy(&th->affine);
    438:   VecDestroyVecs(ts->numcost,&th->VecsDeltaLam);
    439:   VecDestroyVecs(ts->numcost,&th->VecsDeltaMu);
    440:   VecDestroyVecs(ts->numcost,&th->VecsSensiTemp);
    441:   return(0);
    442: }

    446: static PetscErrorCode TSDestroy_Theta(TS ts)
    447: {

    451:   TSReset_Theta(ts);
    452:   PetscFree(ts->data);
    453:   PetscObjectComposeFunction((PetscObject)ts,"TSThetaGetTheta_C",NULL);
    454:   PetscObjectComposeFunction((PetscObject)ts,"TSThetaSetTheta_C",NULL);
    455:   PetscObjectComposeFunction((PetscObject)ts,"TSThetaGetEndpoint_C",NULL);
    456:   PetscObjectComposeFunction((PetscObject)ts,"TSThetaSetEndpoint_C",NULL);
    457:   return(0);
    458: }

    460: /*
    461:   This defines the nonlinear equation that is to be solved with SNES
    462:   G(U) = F[t0+Theta*dt, U, (U-U0)*shift] = 0
    463: */
    466: static PetscErrorCode SNESTSFormFunction_Theta(SNES snes,Vec x,Vec y,TS ts)
    467: {
    468:   TS_Theta       *th = (TS_Theta*)ts->data;
    470:   Vec            X0,Xdot;
    471:   DM             dm,dmsave;
    472:   PetscReal      shift = 1./(th->Theta*ts->time_step);

    475:   SNESGetDM(snes,&dm);
    476:   /* When using the endpoint variant, this is actually 1/Theta * Xdot */
    477:   TSThetaGetX0AndXdot(ts,dm,&X0,&Xdot);
    478:   VecAXPBYPCZ(Xdot,-shift,shift,0,X0,x);

    480:   /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */
    481:   dmsave = ts->dm;
    482:   ts->dm = dm;
    483:   TSComputeIFunction(ts,th->stage_time,x,Xdot,y,PETSC_FALSE);
    484:   ts->dm = dmsave;
    485:   TSThetaRestoreX0AndXdot(ts,dm,&X0,&Xdot);
    486:   return(0);
    487: }

    491: static PetscErrorCode SNESTSFormJacobian_Theta(SNES snes,Vec x,Mat A,Mat B,TS ts)
    492: {
    493:   TS_Theta       *th = (TS_Theta*)ts->data;
    495:   Vec            Xdot;
    496:   DM             dm,dmsave;
    497:   PetscReal      shift = 1./(th->Theta*ts->time_step);

    500:   SNESGetDM(snes,&dm);

    502:   /* th->Xdot has already been computed in SNESTSFormFunction_Theta (SNES guarantees this) */
    503:   TSThetaGetX0AndXdot(ts,dm,NULL,&Xdot);

    505:   dmsave = ts->dm;
    506:   ts->dm = dm;
    507:   TSComputeIJacobian(ts,th->stage_time,x,Xdot,shift,A,B,PETSC_FALSE);
    508:   ts->dm = dmsave;
    509:   TSThetaRestoreX0AndXdot(ts,dm,NULL,&Xdot);
    510:   return(0);
    511: }

    515: static PetscErrorCode TSSetUp_Theta(TS ts)
    516: {
    517:   TS_Theta       *th = (TS_Theta*)ts->data;
    519:   SNES           snes;
    520:   TSAdapt        adapt;
    521:   DM             dm;

    524:   if (!th->X) {
    525:     VecDuplicate(ts->vec_sol,&th->X);
    526:   }
    527:   if (!th->Xdot) {
    528:     VecDuplicate(ts->vec_sol,&th->Xdot);
    529:   }
    530:   if (!th->X0) {
    531:     VecDuplicate(ts->vec_sol,&th->X0);
    532:   }
    533:   TSGetSNES(ts,&snes);
    534:   TSGetDM(ts,&dm);
    535:   if (dm) {
    536:     DMCoarsenHookAdd(dm,DMCoarsenHook_TSTheta,DMRestrictHook_TSTheta,ts);
    537:     DMSubDomainHookAdd(dm,DMSubDomainHook_TSTheta,DMSubDomainRestrictHook_TSTheta,ts);
    538:   }
    539:   if (th->Theta == 0.5 && th->endpoint) th->order = 2;
    540:   else th->order = 1;

    542:   TSGetAdapt(ts,&adapt);
    543:   if (!th->adapt) {
    544:     TSAdaptSetType(adapt,TSADAPTNONE);
    545:   }
    546:   return(0);
    547: }

    551: static PetscErrorCode TSSetUp_BEuler(TS ts)
    552: {
    553:   TS_Theta       *th = (TS_Theta*)ts->data;

    557:   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");
    558:   TSSetUp_Theta(ts);
    559:   return(0);
    560: }

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

    570:   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");
    571:   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");
    572:   TSSetUp_Theta(ts);
    573:   return(0);
    574: }
    575: /*------------------------------------------------------------*/

    579: static PetscErrorCode TSAdjointSetUp_Theta(TS ts)
    580: {
    581:   TS_Theta       *th = (TS_Theta*)ts->data;

    585:   VecDuplicateVecs(ts->vecs_sensi[0],ts->numcost,&th->VecsDeltaLam);
    586:   if(ts->vecs_sensip) {
    587:     VecDuplicateVecs(ts->vecs_sensip[0],ts->numcost,&th->VecsDeltaMu);
    588:   }
    589:   VecDuplicateVecs(ts->vecs_sensi[0],ts->numcost,&th->VecsSensiTemp);
    590:   return(0);
    591: }
    592: /*------------------------------------------------------------*/

    596: static PetscErrorCode TSSetFromOptions_Theta(PetscOptions *PetscOptionsObject,TS ts)
    597: {
    598:   TS_Theta       *th = (TS_Theta*)ts->data;

    602:   PetscOptionsHead(PetscOptionsObject,"Theta ODE solver options");
    603:   {
    604:     PetscOptionsReal("-ts_theta_theta","Location of stage (0<Theta<=1)","TSThetaSetTheta",th->Theta,&th->Theta,NULL);
    605:     PetscOptionsBool("-ts_theta_extrapolate","Extrapolate stage solution from previous solution (sometimes unstable)","TSThetaSetExtrapolate",th->extrapolate,&th->extrapolate,NULL);
    606:     PetscOptionsBool("-ts_theta_endpoint","Use the endpoint instead of midpoint form of the Theta method","TSThetaSetEndpoint",th->endpoint,&th->endpoint,NULL);
    607:     PetscOptionsBool("-ts_theta_adapt","Use time-step adaptivity with the Theta method","",th->adapt,&th->adapt,NULL);
    608:     SNESSetFromOptions(ts->snes);
    609:   }
    610:   PetscOptionsTail();
    611:   return(0);
    612: }

    616: static PetscErrorCode TSView_Theta(TS ts,PetscViewer viewer)
    617: {
    618:   TS_Theta       *th = (TS_Theta*)ts->data;
    619:   PetscBool      iascii;

    623:   PetscObjectTypeCompare((PetscObject)viewer,PETSCVIEWERASCII,&iascii);
    624:   if (iascii) {
    625:     PetscViewerASCIIPrintf(viewer,"  Theta=%g\n",(double)th->Theta);
    626:     PetscViewerASCIIPrintf(viewer,"  Extrapolation=%s\n",th->extrapolate ? "yes" : "no");
    627:   }
    628:   if (ts->snes) {SNESView(ts->snes,viewer);}
    629:   return(0);
    630: }

    634: PetscErrorCode  TSThetaGetTheta_Theta(TS ts,PetscReal *theta)
    635: {
    636:   TS_Theta *th = (TS_Theta*)ts->data;

    639:   *theta = th->Theta;
    640:   return(0);
    641: }

    645: PetscErrorCode  TSThetaSetTheta_Theta(TS ts,PetscReal theta)
    646: {
    647:   TS_Theta *th = (TS_Theta*)ts->data;

    650:   if (theta <= 0 || 1 < theta) SETERRQ1(PetscObjectComm((PetscObject)ts),PETSC_ERR_ARG_OUTOFRANGE,"Theta %g not in range (0,1]",(double)theta);
    651:   th->Theta = theta;
    652:   return(0);
    653: }

    657: PetscErrorCode  TSThetaGetEndpoint_Theta(TS ts,PetscBool *endpoint)
    658: {
    659:   TS_Theta *th = (TS_Theta*)ts->data;

    662:   *endpoint = th->endpoint;
    663:   return(0);
    664: }

    668: PetscErrorCode  TSThetaSetEndpoint_Theta(TS ts,PetscBool flg)
    669: {
    670:   TS_Theta *th = (TS_Theta*)ts->data;

    673:   th->endpoint = flg;
    674:   return(0);
    675: }

    677: #if defined(PETSC_HAVE_COMPLEX)
    680: static PetscErrorCode TSComputeLinearStability_Theta(TS ts,PetscReal xr,PetscReal xi,PetscReal *yr,PetscReal *yi)
    681: {
    682:   PetscComplex z   = xr + xi*PETSC_i,f;
    683:   TS_Theta     *th = (TS_Theta*)ts->data;
    684:   const PetscReal one = 1.0;

    687:   f   = (one + (one - th->Theta)*z)/(one - th->Theta*z);
    688:   *yr = PetscRealPartComplex(f);
    689:   *yi = PetscImaginaryPartComplex(f);
    690:   return(0);
    691: }
    692: #endif

    696: static PetscErrorCode  TSGetStages_Theta(TS ts,PetscInt *ns,Vec **Y)
    697: {
    698:   TS_Theta     *th = (TS_Theta*)ts->data;

    701:   *ns = 1;
    702:   if(Y) {
    703:     *Y  = (th->endpoint)?&(th->X0):&(th->X);
    704:   }
    705:   return(0);
    706: }

    708: /* ------------------------------------------------------------ */
    709: /*MC
    710:       TSTHETA - DAE solver using the implicit Theta method

    712:    Level: beginner

    714:    Options Database:
    715:       -ts_theta_theta <Theta> - Location of stage (0<Theta<=1)
    716:       -ts_theta_extrapolate <flg> Extrapolate stage solution from previous solution (sometimes unstable)
    717:       -ts_theta_endpoint <flag> - Use the endpoint (like Crank-Nicholson) instead of midpoint form of the Theta method

    719:    Notes:
    720: $  -ts_type theta -ts_theta_theta 1.0 corresponds to backward Euler (TSBEULER)
    721: $  -ts_type theta -ts_theta_theta 0.5 corresponds to the implicit midpoint rule
    722: $  -ts_type theta -ts_theta_theta 0.5 -ts_theta_endpoint corresponds to Crank-Nicholson (TSCN)

    726:    This method can be applied to DAE.

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

    730: .vb
    731:   Theta | Theta
    732:   -------------
    733:         |  1
    734: .ve

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

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

    740: .vb
    741:   0 | 0         0
    742:   1 | 1-Theta   Theta
    743:   -------------------
    744:     | 1-Theta   Theta
    745: .ve

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

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

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

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

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

    757: M*/
    760: PETSC_EXTERN PetscErrorCode TSCreate_Theta(TS ts)
    761: {
    762:   TS_Theta       *th;

    766:   ts->ops->reset           = TSReset_Theta;
    767:   ts->ops->destroy         = TSDestroy_Theta;
    768:   ts->ops->view            = TSView_Theta;
    769:   ts->ops->setup           = TSSetUp_Theta;
    770:   ts->ops->adjointsetup    = TSAdjointSetUp_Theta;
    771:   ts->ops->step            = TSStep_Theta;
    772:   ts->ops->interpolate     = TSInterpolate_Theta;
    773:   ts->ops->evaluatestep    = TSEvaluateStep_Theta;
    774:   ts->ops->rollback        = TSRollBack_Theta;
    775:   ts->ops->setfromoptions  = TSSetFromOptions_Theta;
    776:   ts->ops->snesfunction    = SNESTSFormFunction_Theta;
    777:   ts->ops->snesjacobian    = SNESTSFormJacobian_Theta;
    778: #if defined(PETSC_HAVE_COMPLEX)
    779:   ts->ops->linearstability = TSComputeLinearStability_Theta;
    780: #endif
    781:   ts->ops->getstages       = TSGetStages_Theta;
    782:   ts->ops->adjointstep     = TSAdjointStep_Theta;

    784:   PetscNewLog(ts,&th);
    785:   ts->data = (void*)th;

    787:   th->extrapolate = PETSC_FALSE;
    788:   th->Theta       = 0.5;
    789:   th->ccfl        = 1.0;
    790:   th->adapt       = PETSC_FALSE;
    791:   PetscObjectComposeFunction((PetscObject)ts,"TSThetaGetTheta_C",TSThetaGetTheta_Theta);
    792:   PetscObjectComposeFunction((PetscObject)ts,"TSThetaSetTheta_C",TSThetaSetTheta_Theta);
    793:   PetscObjectComposeFunction((PetscObject)ts,"TSThetaGetEndpoint_C",TSThetaGetEndpoint_Theta);
    794:   PetscObjectComposeFunction((PetscObject)ts,"TSThetaSetEndpoint_C",TSThetaSetEndpoint_Theta);
    795:   return(0);
    796: }

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

    803:   Not Collective

    805:   Input Parameter:
    806: .  ts - timestepping context

    808:   Output Parameter:
    809: .  theta - stage abscissa

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

    814:   Level: Advanced

    816: .seealso: TSThetaSetTheta()
    817: @*/
    818: PetscErrorCode  TSThetaGetTheta(TS ts,PetscReal *theta)
    819: {

    825:   PetscUseMethod(ts,"TSThetaGetTheta_C",(TS,PetscReal*),(ts,theta));
    826:   return(0);
    827: }

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

    834:   Not Collective

    836:   Input Parameter:
    837: +  ts - timestepping context
    838: -  theta - stage abscissa

    840:   Options Database:
    841: .  -ts_theta_theta <theta>

    843:   Level: Intermediate

    845: .seealso: TSThetaGetTheta()
    846: @*/
    847: PetscErrorCode  TSThetaSetTheta(TS ts,PetscReal theta)
    848: {

    853:   PetscTryMethod(ts,"TSThetaSetTheta_C",(TS,PetscReal),(ts,theta));
    854:   return(0);
    855: }

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

    862:   Not Collective

    864:   Input Parameter:
    865: .  ts - timestepping context

    867:   Output Parameter:
    868: .  endpoint - PETSC_TRUE when using the endpoint variant

    870:   Level: Advanced

    872: .seealso: TSThetaSetEndpoint(), TSTHETA, TSCN
    873: @*/
    874: PetscErrorCode TSThetaGetEndpoint(TS ts,PetscBool *endpoint)
    875: {

    881:   PetscTryMethod(ts,"TSThetaGetEndpoint_C",(TS,PetscBool*),(ts,endpoint));
    882:   return(0);
    883: }

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

    890:   Not Collective

    892:   Input Parameter:
    893: +  ts - timestepping context
    894: -  flg - PETSC_TRUE to use the endpoint variant

    896:   Options Database:
    897: .  -ts_theta_endpoint <flg>

    899:   Level: Intermediate

    901: .seealso: TSTHETA, TSCN
    902: @*/
    903: PetscErrorCode TSThetaSetEndpoint(TS ts,PetscBool flg)
    904: {

    909:   PetscTryMethod(ts,"TSThetaSetEndpoint_C",(TS,PetscBool),(ts,flg));
    910:   return(0);
    911: }

    913: /*
    914:  * TSBEULER and TSCN are straightforward specializations of TSTHETA.
    915:  * The creation functions for these specializations are below.
    916:  */

    920: static PetscErrorCode TSView_BEuler(TS ts,PetscViewer viewer)
    921: {

    925:   SNESView(ts->snes,viewer);
    926:   return(0);
    927: }

    929: /*MC
    930:       TSBEULER - ODE solver using the implicit backward Euler method

    932:   Level: beginner

    934:   Notes:
    935:   TSBEULER is equivalent to TSTHETA with Theta=1.0

    937: $  -ts_type theta -ts_theta_theta 1.

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

    941: M*/
    944: PETSC_EXTERN PetscErrorCode TSCreate_BEuler(TS ts)
    945: {

    949:   TSCreate_Theta(ts);
    950:   TSThetaSetTheta(ts,1.0);
    951:   ts->ops->setup = TSSetUp_BEuler;
    952:   ts->ops->view = TSView_BEuler;
    953:   return(0);
    954: }

    958: static PetscErrorCode TSView_CN(TS ts,PetscViewer viewer)
    959: {

    963:   SNESView(ts->snes,viewer);
    964:   return(0);
    965: }

    967: /*MC
    968:       TSCN - ODE solver using the implicit Crank-Nicolson method.

    970:   Level: beginner

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

    975: $  -ts_type theta -ts_theta_theta 0.5 -ts_theta_endpoint

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

    979: M*/
    982: PETSC_EXTERN PetscErrorCode TSCreate_CN(TS ts)
    983: {

    987:   TSCreate_Theta(ts);
    988:   TSThetaSetTheta(ts,0.5);
    989:   TSThetaSetEndpoint(ts,PETSC_TRUE);
    990:   ts->ops->setup = TSSetUp_CN;
    991:   ts->ops->view = TSView_CN;
    992:   return(0);
    993: }