Actual source code: theta.c
petsc-3.7.1 2016-05-15
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:    PetscBool    adapt;                    /* Use time-step adaptivity ? */
     27:    Vec          vec_sol_prev;
     28:    Vec          vec_lte_work;

     30:    TSStepStatus status;
     31:  } TS_Theta;

     35:  static PetscErrorCode TSThetaGetX0AndXdot(TS ts,DM dm,Vec *X0,Vec *Xdot)
     36:  {
     37:    TS_Theta       *th = (TS_Theta*)ts->data;

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

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

     61:    if (X0) {
     62:      if (dm && dm != ts->dm) {
     63:        DMRestoreNamedGlobalVector(dm,"TSTheta_X0",X0);
     64:      }
     65:    }
     66:    if (Xdot) {
     67:      if (dm && dm != ts->dm) {
     68:        DMRestoreNamedGlobalVector(dm,"TSTheta_Xdot",Xdot);
     69:      }
     70:    }
     71:    return(0);
     72:  }

     76:  static PetscErrorCode DMCoarsenHook_TSTheta(DM fine,DM coarse,void *ctx)
     77:  {

     80:    return(0);
     81:  }

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

     92:    TSThetaGetX0AndXdot(ts,fine,&X0,&Xdot);
     93:    TSThetaGetX0AndXdot(ts,coarse,&X0_c,&Xdot_c);
     94:    MatRestrict(restrct,X0,X0_c);
     95:    MatRestrict(restrct,Xdot,Xdot_c);
     96:    VecPointwiseMult(X0_c,rscale,X0_c);
     97:    VecPointwiseMult(Xdot_c,rscale,Xdot_c);
     98:    TSThetaRestoreX0AndXdot(ts,fine,&X0,&Xdot);
     99:    TSThetaRestoreX0AndXdot(ts,coarse,&X0_c,&Xdot_c);
    100:    return(0);
    101:  }

    105:  static PetscErrorCode DMSubDomainHook_TSTheta(DM dm,DM subdm,void *ctx)
    106:  {

    109:    return(0);
    110:  }

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

    121:    TSThetaGetX0AndXdot(ts,dm,&X0,&Xdot);
    122:    TSThetaGetX0AndXdot(ts,subdm,&X0_sub,&Xdot_sub);

    124:    VecScatterBegin(gscat,X0,X0_sub,INSERT_VALUES,SCATTER_FORWARD);
    125:    VecScatterEnd(gscat,X0,X0_sub,INSERT_VALUES,SCATTER_FORWARD);

    127:    VecScatterBegin(gscat,Xdot,Xdot_sub,INSERT_VALUES,SCATTER_FORWARD);
    128:    VecScatterEnd(gscat,Xdot,Xdot_sub,INSERT_VALUES,SCATTER_FORWARD);

    130:    TSThetaRestoreX0AndXdot(ts,dm,&X0,&Xdot);
    131:    TSThetaRestoreX0AndXdot(ts,subdm,&X0_sub,&Xdot_sub);
    132:    return(0);
    133:  }

    137:  static PetscErrorCode TSForwardCostIntegral_Theta(TS ts)
    138:  {
    139:    TS_Theta       *th = (TS_Theta*)ts->data;

    143:    /* backup cost integral */
    144:    VecCopy(ts->vec_costintegral,th->VecCostIntegral0);
    145:    if (th->endpoint) {
    146:      TSAdjointComputeCostIntegrand(ts,th->ptime,th->X0,ts->vec_costintegrand);
    147:      VecAXPY(ts->vec_costintegral,th->time_step*(1-th->Theta),ts->vec_costintegrand);
    148:    }
    149:    TSAdjointComputeCostIntegrand(ts,th->stage_time,th->X,ts->vec_costintegrand);
    150:    if (th->endpoint) {
    151:      VecAXPY(ts->vec_costintegral,th->time_step*th->Theta,ts->vec_costintegrand);
    152:    } else {
    153:      VecAXPY(ts->vec_costintegral,th->time_step,ts->vec_costintegrand);
    154:    }
    155:    return(0);
    156:  }

    160:  static PetscErrorCode TSAdjointCostIntegral_Theta(TS ts)
    161:  {
    162:    TS_Theta       *th = (TS_Theta*)ts->data;

    166:    if (th->endpoint) {
    167:      /* Evolve ts->vec_costintegral to compute integrals */
    168:      TSAdjointComputeCostIntegrand(ts,ts->ptime,ts->vec_sol,ts->vec_costintegrand);
    169:      VecAXPY(ts->vec_costintegral,-ts->time_step*th->Theta,ts->vec_costintegrand);
    170:      if (th->Theta!=1) {
    171:        TSAdjointComputeCostIntegrand(ts,th->ptime,th->X0,ts->vec_costintegrand);
    172:        VecAXPY(ts->vec_costintegral,ts->time_step*(th->Theta-1),ts->vec_costintegrand);
    173:      }
    174:    }else {
    175:      TSAdjointComputeCostIntegrand(ts,th->stage_time,th->X,ts->vec_costintegrand);
    176:      VecAXPY(ts->vec_costintegral,-ts->time_step,ts->vec_costintegrand);
    177:    }
    178:    return(0);
    179:  }

    183:  static PetscErrorCode TS_SNESSolve(TS ts,Vec b,Vec x)
    184:  {
    185:    PetscInt       nits,lits;

    189:    SNESSolve(ts->snes,b,x);
    190:    SNESGetIterationNumber(ts->snes,&nits);
    191:    SNESGetLinearSolveIterations(ts->snes,&lits);
    192:    ts->snes_its += nits; ts->ksp_its += lits;
    193:    return(0);
    194:  }

    198:  static PetscErrorCode TSStep_Theta(TS ts)
    199:  {
    200:    TS_Theta       *th = (TS_Theta*)ts->data;
    201:    PetscInt       rejections = 0;
    202:    PetscBool      stageok,accept = PETSC_TRUE;
    203:    PetscReal      next_time_step = ts->time_step;

    207:    if (!ts->steprollback) {
    208:      if (th->adapt) { VecCopy(th->X0,th->vec_sol_prev); }
    209:      VecCopy(ts->vec_sol,th->X0);
    210:    }

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

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

    218:      VecCopy(th->X0,th->X);
    219:      if (th->extrapolate && !ts->steprestart) {
    220:        VecAXPY(th->X,1/shift,th->Xdot);
    221:      }
    222:      if (th->endpoint) { /* This formulation assumes linear time-independent mass matrix */
    223:        if (!th->affine) {VecDuplicate(ts->vec_sol,&th->affine);}
    224:        VecZeroEntries(th->Xdot);
    225:        TSComputeIFunction(ts,ts->ptime,th->X0,th->Xdot,th->affine,PETSC_FALSE);
    226:        VecScale(th->affine,(th->Theta-1)/th->Theta);
    227:      } else if (th->affine) { /* Just in case th->endpoint is changed between calls to TSStep_Theta() */
    228:        VecZeroEntries(th->affine);
    229:      }
    230:      TSPreStage(ts,th->stage_time);
    231:      TS_SNESSolve(ts,th->affine,th->X);
    232:      TSPostStage(ts,th->stage_time,0,&th->X);
    233:      TSAdaptCheckStage(ts->adapt,ts,th->stage_time,th->X,&stageok);
    234:      if (!stageok) goto reject_step;

    236:      th->status = TS_STEP_PENDING;
    237:      if (th->endpoint) {
    238:        VecCopy(th->X,ts->vec_sol);
    239:      } else {
    240:        VecAXPBYPCZ(th->Xdot,-shift,shift,0,th->X0,th->X);
    241:        VecAXPY(ts->vec_sol,ts->time_step,th->Xdot);
    242:      }
    243:      TSAdaptChoose(ts->adapt,ts,ts->time_step,NULL,&next_time_step,&accept);
    244:      th->status = accept ? TS_STEP_COMPLETE : TS_STEP_INCOMPLETE;
    245:      if (!accept) {
    246:        VecCopy(th->X0,ts->vec_sol);
    247:        ts->time_step = next_time_step;
    248:        goto reject_step;
    249:      }

    251:      if (ts->costintegralfwd) { /* Save the info for the later use in cost integral evaluation */
    252:        th->ptime     = ts->ptime;
    253:        th->time_step = ts->time_step;
    254:      }

    256:      ts->ptime += ts->time_step;
    257:      ts->time_step = next_time_step;
    258:      break;

    260:    reject_step:
    261:      ts->reject++; accept = PETSC_FALSE;
    262:      if (!ts->reason && ++rejections > ts->max_reject && ts->max_reject >= 0) {
    263:        ts->reason = TS_DIVERGED_STEP_REJECTED;
    264:        PetscInfo2(ts,"Step=%D, step rejections %D greater than current TS allowed, stopping solve\n",ts->steps,rejections);
    265:      }
    266:    }
    267:    return(0);
    268:  }

    272:  static PetscErrorCode TSAdjointStep_Theta(TS ts)
    273:  {
    274:    TS_Theta            *th = (TS_Theta*)ts->data;
    275:    Vec                 *VecsDeltaLam = th->VecsDeltaLam,*VecsDeltaMu = th->VecsDeltaMu,*VecsSensiTemp = th->VecsSensiTemp;
    276:    PetscInt            nadj;
    277:    PetscErrorCode      ierr;
    278:    Mat                 J,Jp;
    279:    KSP                 ksp;
    280:    PetscReal           shift;

    284:    th->status = TS_STEP_INCOMPLETE;
    285:    SNESGetKSP(ts->snes,&ksp);
    286:    TSGetIJacobian(ts,&J,&Jp,NULL,NULL);

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

    292:    /* Build RHS */
    293:    if (ts->vec_costintegral) { /* Cost function has an integral term */
    294:      if (th->endpoint) {
    295:        TSAdjointComputeDRDYFunction(ts,ts->ptime,ts->vec_sol,ts->vecs_drdy);
    296:      }else {
    297:        TSAdjointComputeDRDYFunction(ts,th->stage_time,th->X,ts->vecs_drdy);
    298:      }
    299:    }
    300:    for (nadj=0; nadj<ts->numcost; nadj++) {
    301:      VecCopy(ts->vecs_sensi[nadj],VecsSensiTemp[nadj]);
    302:      VecScale(VecsSensiTemp[nadj],-1./(th->Theta*ts->time_step));
    303:      if (ts->vec_costintegral) {
    304:        VecAXPY(VecsSensiTemp[nadj],1.,ts->vecs_drdy[nadj]);
    305:      }
    306:    }

    308:    /* Build LHS */
    309:    shift = -1./(th->Theta*ts->time_step);
    310:    if (th->endpoint) {
    311:      TSComputeIJacobian(ts,ts->ptime,ts->vec_sol,th->Xdot,shift,J,Jp,PETSC_FALSE);
    312:    }else {
    313:      TSComputeIJacobian(ts,th->stage_time,th->X,th->Xdot,shift,J,Jp,PETSC_FALSE);
    314:    }
    315:    KSPSetOperators(ksp,J,Jp);

    317:    /* Solve LHS X = RHS */
    318:    for (nadj=0; nadj<ts->numcost; nadj++) {
    319:      KSPSolveTranspose(ksp,VecsSensiTemp[nadj],VecsDeltaLam[nadj]);
    320:    }

    322:    /* Update sensitivities, and evaluate integrals if there is any */
    323:    if(th->endpoint) { /* two-stage case */
    324:      if (th->Theta!=1.) {
    325:        shift = -1./((th->Theta-1.)*ts->time_step);
    326:        TSComputeIJacobian(ts,th->ptime,th->X0,th->Xdot,shift,J,Jp,PETSC_FALSE);
    327:        if (ts->vec_costintegral) {
    328:          TSAdjointComputeDRDYFunction(ts,th->ptime,th->X0,ts->vecs_drdy);
    329:        }
    330:        for (nadj=0; nadj<ts->numcost; nadj++) {
    331:          MatMultTranspose(J,VecsDeltaLam[nadj],ts->vecs_sensi[nadj]);
    332:          if (ts->vec_costintegral) {
    333:            VecAXPY(ts->vecs_sensi[nadj],-1.,ts->vecs_drdy[nadj]);
    334:          }
    335:          VecScale(ts->vecs_sensi[nadj],1./shift);
    336:        }
    337:      }else { /* backward Euler */
    338:        shift = 0.0;
    339:        TSComputeIJacobian(ts,ts->ptime,ts->vec_sol,th->Xdot,shift,J,Jp,PETSC_FALSE); /* get -f_y */
    340:        for (nadj=0; nadj<ts->numcost; nadj++) {
    341:          MatMultTranspose(J,VecsDeltaLam[nadj],VecsSensiTemp[nadj]);
    342:          VecAXPY(ts->vecs_sensi[nadj],ts->time_step,VecsSensiTemp[nadj]);
    343:          if (ts->vec_costintegral) {
    344:            VecAXPY(ts->vecs_sensi[nadj],-ts->time_step,ts->vecs_drdy[nadj]);
    345:          }
    346:        }
    347:      }

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

    403:    th->status = TS_STEP_COMPLETE;
    404:    return(0);
    405:  }

    409:  static PetscErrorCode TSInterpolate_Theta(TS ts,PetscReal t,Vec X)
    410:  {
    411:    TS_Theta       *th = (TS_Theta*)ts->data;
    412:    PetscReal      dt  = t - ts->ptime;

    416:    VecCopy(ts->vec_sol,th->X);
    417:    if (th->endpoint) dt *= th->Theta;
    418:    VecWAXPY(X,dt,th->Xdot,th->X);
    419:    return(0);
    420:  }

    424: static PetscErrorCode TSEvaluateWLTE_Theta(TS ts,NormType wnormtype,PetscInt *order,PetscReal *wlte)
    425: {
    426:   TS_Theta       *th = (TS_Theta*)ts->data;
    427:   Vec            X = ts->vec_sol;      /* X = solution */
    428:   Vec            Y = th->vec_lte_work; /* Y = X + LTE  */

    432:   /* Cannot compute LTE in first step or in restart after event */
    433:   if (ts->steprestart) {*wlte = -1; return(0);}
    434:   /* Compute LTE using backward differences with non-constant time step */
    435:   {
    436:     PetscReal   h = ts->time_step, h_prev = ts->ptime - ts->ptime_prev;
    437:     PetscReal   a = 1 + h_prev/h;
    438:     PetscScalar scal[3]; Vec vecs[3];
    439:     scal[0] = +1/a; scal[1] = -1/(a-1); scal[2] = +1/(a*(a-1));
    440:     vecs[0] = X;    vecs[1] = th->X0;   vecs[2] = th->vec_sol_prev;
    441:     VecCopy(X,Y);
    442:     VecMAXPY(Y,3,scal,vecs);
    443:     TSErrorWeightedNorm(ts,X,Y,wnormtype,wlte);
    444:   }
    445:   if (order) *order = 2;
    446:   return(0);
    447: }

    451: static PetscErrorCode TSRollBack_Theta(TS ts)
    452: {
    453:   TS_Theta       *th = (TS_Theta*)ts->data;

    457:   VecCopy(th->X0,ts->vec_sol);
    458:   if (ts->vec_costintegral && ts->costintegralfwd) {
    459:     VecCopy(th->VecCostIntegral0,ts->vec_costintegral);
    460:   }
    461:   return(0);
    462: }

    464: /*------------------------------------------------------------*/
    467: static PetscErrorCode TSReset_Theta(TS ts)
    468: {
    469:   TS_Theta       *th = (TS_Theta*)ts->data;

    473:   VecDestroy(&th->X);
    474:   VecDestroy(&th->Xdot);
    475:   VecDestroy(&th->X0);
    476:   VecDestroy(&th->affine);

    478:   VecDestroy(&th->vec_sol_prev);
    479:   VecDestroy(&th->vec_lte_work);

    481:   VecDestroy(&th->VecCostIntegral0);
    482:   VecDestroyVecs(ts->numcost,&th->VecsDeltaLam);
    483:   VecDestroyVecs(ts->numcost,&th->VecsDeltaMu);
    484:   VecDestroyVecs(ts->numcost,&th->VecsSensiTemp);
    485:   return(0);
    486: }

    490: static PetscErrorCode TSDestroy_Theta(TS ts)
    491: {

    495:   TSReset_Theta(ts);
    496:   PetscFree(ts->data);
    497:   PetscObjectComposeFunction((PetscObject)ts,"TSThetaGetTheta_C",NULL);
    498:   PetscObjectComposeFunction((PetscObject)ts,"TSThetaSetTheta_C",NULL);
    499:   PetscObjectComposeFunction((PetscObject)ts,"TSThetaGetEndpoint_C",NULL);
    500:   PetscObjectComposeFunction((PetscObject)ts,"TSThetaSetEndpoint_C",NULL);
    501:   return(0);
    502: }

    504: /*
    505:   This defines the nonlinear equation that is to be solved with SNES
    506:   G(U) = F[t0+Theta*dt, U, (U-U0)*shift] = 0
    507: */
    510: static PetscErrorCode SNESTSFormFunction_Theta(SNES snes,Vec x,Vec y,TS ts)
    511: {
    512:   TS_Theta       *th = (TS_Theta*)ts->data;
    514:   Vec            X0,Xdot;
    515:   DM             dm,dmsave;
    516:   PetscReal      shift = 1/(th->Theta*ts->time_step);

    519:   SNESGetDM(snes,&dm);
    520:   /* When using the endpoint variant, this is actually 1/Theta * Xdot */
    521:   TSThetaGetX0AndXdot(ts,dm,&X0,&Xdot);
    522:   VecAXPBYPCZ(Xdot,-shift,shift,0,X0,x);

    524:   /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */
    525:   dmsave = ts->dm;
    526:   ts->dm = dm;
    527:   TSComputeIFunction(ts,th->stage_time,x,Xdot,y,PETSC_FALSE);
    528:   ts->dm = dmsave;
    529:   TSThetaRestoreX0AndXdot(ts,dm,&X0,&Xdot);
    530:   return(0);
    531: }

    535: static PetscErrorCode SNESTSFormJacobian_Theta(SNES snes,Vec x,Mat A,Mat B,TS ts)
    536: {
    537:   TS_Theta       *th = (TS_Theta*)ts->data;
    539:   Vec            Xdot;
    540:   DM             dm,dmsave;
    541:   PetscReal      shift = 1/(th->Theta*ts->time_step);

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

    548:   dmsave = ts->dm;
    549:   ts->dm = dm;
    550:   TSComputeIJacobian(ts,th->stage_time,x,Xdot,shift,A,B,PETSC_FALSE);
    551:   ts->dm = dmsave;
    552:   TSThetaRestoreX0AndXdot(ts,dm,NULL,&Xdot);
    553:   return(0);
    554: }

    558: static PetscErrorCode TSSetUp_Theta(TS ts)
    559: {
    560:   TS_Theta       *th = (TS_Theta*)ts->data;

    564:   if (!th->VecCostIntegral0 && ts->vec_costintegral && ts->costintegralfwd) { /* back up cost integral */
    565:     VecDuplicate(ts->vec_costintegral,&th->VecCostIntegral0);
    566:   }
    567:   if (!th->X) {
    568:     VecDuplicate(ts->vec_sol,&th->X);
    569:   }
    570:   if (!th->Xdot) {
    571:     VecDuplicate(ts->vec_sol,&th->Xdot);
    572:   }
    573:   if (!th->X0) {
    574:     VecDuplicate(ts->vec_sol,&th->X0);
    575:   }
    576:   if (th->endpoint) {
    577:     VecDuplicate(ts->vec_sol,&th->affine);
    578:   }

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

    582:   TSGetDM(ts,&ts->dm);
    583:   DMCoarsenHookAdd(ts->dm,DMCoarsenHook_TSTheta,DMRestrictHook_TSTheta,ts);
    584:   DMSubDomainHookAdd(ts->dm,DMSubDomainHook_TSTheta,DMSubDomainRestrictHook_TSTheta,ts);

    586:   TSGetAdapt(ts,&ts->adapt);
    587:   TSAdaptCandidatesClear(ts->adapt);
    588:   if (!th->adapt) {
    589:     TSAdaptSetType(ts->adapt,TSADAPTNONE);
    590:   } else {
    591:     VecDuplicate(ts->vec_sol,&th->vec_sol_prev);
    592:     VecDuplicate(ts->vec_sol,&th->vec_lte_work);
    593:     if (ts->exact_final_time == TS_EXACTFINALTIME_UNSPECIFIED)
    594:       ts->exact_final_time = TS_EXACTFINALTIME_MATCHSTEP;
    595:   }

    597:   TSGetSNES(ts,&ts->snes);
    598:   return(0);
    599: }

    601: /*------------------------------------------------------------*/

    605: static PetscErrorCode TSAdjointSetUp_Theta(TS ts)
    606: {
    607:   TS_Theta       *th = (TS_Theta*)ts->data;

    611:   VecDuplicateVecs(ts->vecs_sensi[0],ts->numcost,&th->VecsDeltaLam);
    612:   if(ts->vecs_sensip) {
    613:     VecDuplicateVecs(ts->vecs_sensip[0],ts->numcost,&th->VecsDeltaMu);
    614:   }
    615:   VecDuplicateVecs(ts->vecs_sensi[0],ts->numcost,&th->VecsSensiTemp);
    616:   return(0);
    617: }
    618: /*------------------------------------------------------------*/

    622: static PetscErrorCode TSSetFromOptions_Theta(PetscOptionItems *PetscOptionsObject,TS ts)
    623: {
    624:   TS_Theta       *th = (TS_Theta*)ts->data;

    628:   PetscOptionsHead(PetscOptionsObject,"Theta ODE solver options");
    629:   {
    630:     PetscOptionsReal("-ts_theta_theta","Location of stage (0<Theta<=1)","TSThetaSetTheta",th->Theta,&th->Theta,NULL);
    631:     PetscOptionsBool("-ts_theta_endpoint","Use the endpoint instead of midpoint form of the Theta method","TSThetaSetEndpoint",th->endpoint,&th->endpoint,NULL);
    632:     PetscOptionsBool("-ts_theta_adapt","Use time-step adaptivity with the Theta method","",th->adapt,&th->adapt,NULL);
    633:     PetscOptionsBool("-ts_theta_initial_guess_extrapolate","Extrapolate stage initial guess from previous solution (sometimes unstable)","TSThetaSetExtrapolate",th->extrapolate,&th->extrapolate,NULL);
    634:   }
    635:   PetscOptionsTail();
    636:   return(0);
    637: }

    641: static PetscErrorCode TSView_Theta(TS ts,PetscViewer viewer)
    642: {
    643:   TS_Theta       *th = (TS_Theta*)ts->data;
    644:   PetscBool      iascii;

    648:   PetscObjectTypeCompare((PetscObject)viewer,PETSCVIEWERASCII,&iascii);
    649:   if (iascii) {
    650:     PetscViewerASCIIPrintf(viewer,"  Theta=%g\n",(double)th->Theta);
    651:     PetscViewerASCIIPrintf(viewer,"  Extrapolation=%s\n",th->extrapolate ? "yes" : "no");
    652:   }
    653:   if (ts->adapt) {TSAdaptView(ts->adapt,viewer);}
    654:   if (ts->snes)  {SNESView(ts->snes,viewer);}
    655:   return(0);
    656: }

    660: static PetscErrorCode TSThetaGetTheta_Theta(TS ts,PetscReal *theta)
    661: {
    662:   TS_Theta *th = (TS_Theta*)ts->data;

    665:   *theta = th->Theta;
    666:   return(0);
    667: }

    671: static PetscErrorCode TSThetaSetTheta_Theta(TS ts,PetscReal theta)
    672: {
    673:   TS_Theta *th = (TS_Theta*)ts->data;

    676:   if (theta <= 0 || 1 < theta) SETERRQ1(PetscObjectComm((PetscObject)ts),PETSC_ERR_ARG_OUTOFRANGE,"Theta %g not in range (0,1]",(double)theta);
    677:   th->Theta = theta;
    678:   th->order = (th->Theta == 0.5) ? 2 : 1;
    679:   return(0);
    680: }

    684: static PetscErrorCode TSThetaGetEndpoint_Theta(TS ts,PetscBool *endpoint)
    685: {
    686:   TS_Theta *th = (TS_Theta*)ts->data;

    689:   *endpoint = th->endpoint;
    690:   return(0);
    691: }

    695: static PetscErrorCode TSThetaSetEndpoint_Theta(TS ts,PetscBool flg)
    696: {
    697:   TS_Theta *th = (TS_Theta*)ts->data;

    700:   th->endpoint = flg;
    701:   return(0);
    702: }

    704: #if defined(PETSC_HAVE_COMPLEX)
    707: static PetscErrorCode TSComputeLinearStability_Theta(TS ts,PetscReal xr,PetscReal xi,PetscReal *yr,PetscReal *yi)
    708: {
    709:   PetscComplex   z   = xr + xi*PETSC_i,f;
    710:   TS_Theta       *th = (TS_Theta*)ts->data;
    711:   const PetscReal one = 1.0;

    714:   f   = (one + (one - th->Theta)*z)/(one - th->Theta*z);
    715:   *yr = PetscRealPartComplex(f);
    716:   *yi = PetscImaginaryPartComplex(f);
    717:   return(0);
    718: }
    719: #endif

    723: static PetscErrorCode TSGetStages_Theta(TS ts,PetscInt *ns,Vec **Y)
    724: {
    725:   TS_Theta     *th = (TS_Theta*)ts->data;

    728:   if (ns) *ns = 1;
    729:   if (Y)  *Y  = th->endpoint ? &(th->X0) : &(th->X);
    730:   return(0);
    731: }

    733: /* ------------------------------------------------------------ */
    734: /*MC
    735:       TSTHETA - DAE solver using the implicit Theta method

    737:    Level: beginner

    739:    Options Database:
    740: +  -ts_theta_theta <Theta> - Location of stage (0<Theta<=1)
    741: .  -ts_theta_endpoint <flag> - Use the endpoint (like Crank-Nicholson) instead of midpoint form of the Theta method
    742: .  -ts_theta_adapt <flg> - Use time-step adaptivity with the Theta method
    743: -  -ts_theta_initial_guess_extrapolate <flg> - Extrapolate stage initial guess from previous solution (sometimes unstable)

    745:    Notes:
    746: $  -ts_type theta -ts_theta_theta 1.0 corresponds to backward Euler (TSBEULER)
    747: $  -ts_type theta -ts_theta_theta 0.5 corresponds to the implicit midpoint rule
    748: $  -ts_type theta -ts_theta_theta 0.5 -ts_theta_endpoint corresponds to Crank-Nicholson (TSCN)

    750:    This method can be applied to DAE.

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

    754: .vb
    755:   Theta | Theta
    756:   -------------
    757:         |  1
    758: .ve

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

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

    764: .vb
    765:   0 | 0         0
    766:   1 | 1-Theta   Theta
    767:   -------------------
    768:     | 1-Theta   Theta
    769: .ve

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

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

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

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

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

    781: M*/
    784: PETSC_EXTERN PetscErrorCode TSCreate_Theta(TS ts)
    785: {
    786:   TS_Theta       *th;

    790:   ts->ops->reset           = TSReset_Theta;
    791:   ts->ops->destroy         = TSDestroy_Theta;
    792:   ts->ops->view            = TSView_Theta;
    793:   ts->ops->setup           = TSSetUp_Theta;
    794:   ts->ops->adjointsetup    = TSAdjointSetUp_Theta;
    795:   ts->ops->step            = TSStep_Theta;
    796:   ts->ops->interpolate     = TSInterpolate_Theta;
    797:   ts->ops->evaluatewlte    = TSEvaluateWLTE_Theta;
    798:   ts->ops->rollback        = TSRollBack_Theta;
    799:   ts->ops->setfromoptions  = TSSetFromOptions_Theta;
    800:   ts->ops->snesfunction    = SNESTSFormFunction_Theta;
    801:   ts->ops->snesjacobian    = SNESTSFormJacobian_Theta;
    802: #if defined(PETSC_HAVE_COMPLEX)
    803:   ts->ops->linearstability = TSComputeLinearStability_Theta;
    804: #endif
    805:   ts->ops->getstages       = TSGetStages_Theta;
    806:   ts->ops->adjointstep     = TSAdjointStep_Theta;
    807:   ts->ops->adjointintegral = TSAdjointCostIntegral_Theta;
    808:   ts->ops->forwardintegral = TSForwardCostIntegral_Theta;

    810:   PetscNewLog(ts,&th);
    811:   ts->data = (void*)th;

    813:   th->extrapolate = PETSC_FALSE;
    814:   th->Theta       = 0.5;
    815:   th->order       = 2;
    816:   th->adapt       = PETSC_FALSE;
    817:   PetscObjectComposeFunction((PetscObject)ts,"TSThetaGetTheta_C",TSThetaGetTheta_Theta);
    818:   PetscObjectComposeFunction((PetscObject)ts,"TSThetaSetTheta_C",TSThetaSetTheta_Theta);
    819:   PetscObjectComposeFunction((PetscObject)ts,"TSThetaGetEndpoint_C",TSThetaGetEndpoint_Theta);
    820:   PetscObjectComposeFunction((PetscObject)ts,"TSThetaSetEndpoint_C",TSThetaSetEndpoint_Theta);
    821:   return(0);
    822: }

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

    829:   Not Collective

    831:   Input Parameter:
    832: .  ts - timestepping context

    834:   Output Parameter:
    835: .  theta - stage abscissa

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

    840:   Level: Advanced

    842: .seealso: TSThetaSetTheta()
    843: @*/
    844: PetscErrorCode  TSThetaGetTheta(TS ts,PetscReal *theta)
    845: {

    851:   PetscUseMethod(ts,"TSThetaGetTheta_C",(TS,PetscReal*),(ts,theta));
    852:   return(0);
    853: }

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

    860:   Not Collective

    862:   Input Parameter:
    863: +  ts - timestepping context
    864: -  theta - stage abscissa

    866:   Options Database:
    867: .  -ts_theta_theta <theta>

    869:   Level: Intermediate

    871: .seealso: TSThetaGetTheta()
    872: @*/
    873: PetscErrorCode  TSThetaSetTheta(TS ts,PetscReal theta)
    874: {

    879:   PetscTryMethod(ts,"TSThetaSetTheta_C",(TS,PetscReal),(ts,theta));
    880:   return(0);
    881: }

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

    888:   Not Collective

    890:   Input Parameter:
    891: .  ts - timestepping context

    893:   Output Parameter:
    894: .  endpoint - PETSC_TRUE when using the endpoint variant

    896:   Level: Advanced

    898: .seealso: TSThetaSetEndpoint(), TSTHETA, TSCN
    899: @*/
    900: PetscErrorCode TSThetaGetEndpoint(TS ts,PetscBool *endpoint)
    901: {

    907:   PetscUseMethod(ts,"TSThetaGetEndpoint_C",(TS,PetscBool*),(ts,endpoint));
    908:   return(0);
    909: }

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

    916:   Not Collective

    918:   Input Parameter:
    919: +  ts - timestepping context
    920: -  flg - PETSC_TRUE to use the endpoint variant

    922:   Options Database:
    923: .  -ts_theta_endpoint <flg>

    925:   Level: Intermediate

    927: .seealso: TSTHETA, TSCN
    928: @*/
    929: PetscErrorCode TSThetaSetEndpoint(TS ts,PetscBool flg)
    930: {

    935:   PetscTryMethod(ts,"TSThetaSetEndpoint_C",(TS,PetscBool),(ts,flg));
    936:   return(0);
    937: }

    939: /*
    940:  * TSBEULER and TSCN are straightforward specializations of TSTHETA.
    941:  * The creation functions for these specializations are below.
    942:  */

    946: static PetscErrorCode TSSetUp_BEuler(TS ts)
    947: {
    948:   TS_Theta       *th = (TS_Theta*)ts->data;

    952:   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");
    953:   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");
    954:   TSSetUp_Theta(ts);
    955:   return(0);
    956: }

    960: static PetscErrorCode TSView_BEuler(TS ts,PetscViewer viewer)
    961: {

    965:   if (ts->adapt) {TSAdaptView(ts->adapt,viewer);}
    966:   if (ts->snes)  {SNESView(ts->snes,viewer);}
    967:   return(0);
    968: }

    970: /*MC
    971:       TSBEULER - ODE solver using the implicit backward Euler method

    973:   Level: beginner

    975:   Notes:
    976:   TSBEULER is equivalent to TSTHETA with Theta=1.0

    978: $  -ts_type theta -ts_theta_theta 1.0

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

    982: M*/
    985: PETSC_EXTERN PetscErrorCode TSCreate_BEuler(TS ts)
    986: {

    990:   TSCreate_Theta(ts);
    991:   TSThetaSetTheta(ts,1.0);
    992:   TSThetaSetEndpoint(ts,PETSC_FALSE);
    993:   ts->ops->setup = TSSetUp_BEuler;
    994:   ts->ops->view  = TSView_BEuler;
    995:   return(0);
    996: }

    1000: static PetscErrorCode TSSetUp_CN(TS ts)
    1001: {
    1002:   TS_Theta       *th = (TS_Theta*)ts->data;

    1006:   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");
    1007:   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");
    1008:   TSSetUp_Theta(ts);
    1009:   return(0);
    1010: }

    1014: static PetscErrorCode TSView_CN(TS ts,PetscViewer viewer)
    1015: {

    1019:   if (ts->adapt) {TSAdaptView(ts->adapt,viewer);}
    1020:   if (ts->snes)  {SNESView(ts->snes,viewer);}
    1021:   return(0);
    1022: }

    1024: /*MC
    1025:       TSCN - ODE solver using the implicit Crank-Nicolson method.

    1027:   Level: beginner

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

    1032: $  -ts_type theta -ts_theta_theta 0.5 -ts_theta_endpoint

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

    1036: M*/
    1039: PETSC_EXTERN PetscErrorCode TSCreate_CN(TS ts)
    1040: {

    1044:   TSCreate_Theta(ts);
    1045:   TSThetaSetTheta(ts,0.5);
    1046:   TSThetaSetEndpoint(ts,PETSC_TRUE);
    1047:   ts->ops->setup = TSSetUp_CN;
    1048:   ts->ops->view  = TSView_CN;
    1049:   return(0);
    1050: }