Actual source code: theta.c

petsc-master 2019-08-19
Report Typos and Errors
  1: /*
  2:   Code for timestepping with implicit Theta method
  3: */
  4:  #include <petsc/private/tsimpl.h>
  5:  #include <petscsnes.h>
  6:  #include <petscdm.h>
  7:  #include <petscmat.h>

  9: typedef struct {
 10:   /* context for time stepping */
 11:   PetscReal    stage_time;
 12:   Vec          X0,X,Xdot;                /* Storage for stages and time derivative */
 13:   Vec          affine;                   /* Affine vector needed for residual at beginning of step in endpoint formulation */
 14:   PetscReal    Theta;
 15:   PetscReal    ptime;
 16:   PetscReal    time_step;
 17:   PetscInt     order;
 18:   PetscBool    endpoint;
 19:   PetscBool    extrapolate;
 20:   TSStepStatus status;
 21:   Vec          VecCostIntegral0;         /* Backup for roll-backs due to events */

 23:   /* context for sensitivity analysis */
 24:   PetscInt     num_tlm;                  /* Total number of tangent linear equations */
 25:   Vec          *VecsDeltaLam;            /* Increment of the adjoint sensitivity w.r.t IC at stage */
 26:   Vec          *VecsDeltaMu;             /* Increment of the adjoint sensitivity w.r.t P at stage */
 27:   Vec          *VecsSensiTemp;           /* Vector to be multiplied with Jacobian transpose */
 28:   Mat          MatDeltaFwdSensip;        /* Increment of the forward sensitivity at stage */
 29:   Vec          VecDeltaFwdSensipCol;     /* Working vector for holding one column of the sensitivity matrix */
 30:   Mat          MatFwdSensip0;            /* backup for roll-backs due to events */
 31:   Mat          MatIntegralSensipTemp;    /* Working vector for forward integral sensitivity */
 32:   Mat          MatIntegralSensip0;       /* backup for roll-backs due to events */
 33:   Vec          *VecsDeltaLam2;           /* Increment of the 2nd-order adjoint sensitivity w.r.t IC at stage */
 34:   Vec          *VecsDeltaMu2;            /* Increment of the 2nd-order adjoint sensitivity w.r.t P at stage */
 35:   Vec          *VecsSensi2Temp;          /* Working vectors that holds the residual for the second-order adjoint */
 36:   Vec          *VecsAffine;              /* Working vectors to store residuals */
 37:   /* context for error estimation */
 38:   Vec          vec_sol_prev;
 39:   Vec          vec_lte_work;
 40: } TS_Theta;

 42: static PetscErrorCode TSThetaGetX0AndXdot(TS ts,DM dm,Vec *X0,Vec *Xdot)
 43: {
 44:   TS_Theta       *th = (TS_Theta*)ts->data;

 48:   if (X0) {
 49:     if (dm && dm != ts->dm) {
 50:       DMGetNamedGlobalVector(dm,"TSTheta_X0",X0);
 51:     } else *X0 = ts->vec_sol;
 52:   }
 53:   if (Xdot) {
 54:     if (dm && dm != ts->dm) {
 55:       DMGetNamedGlobalVector(dm,"TSTheta_Xdot",Xdot);
 56:     } else *Xdot = th->Xdot;
 57:   }
 58:   return(0);
 59: }

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

 66:   if (X0) {
 67:     if (dm && dm != ts->dm) {
 68:       DMRestoreNamedGlobalVector(dm,"TSTheta_X0",X0);
 69:     }
 70:   }
 71:   if (Xdot) {
 72:     if (dm && dm != ts->dm) {
 73:       DMRestoreNamedGlobalVector(dm,"TSTheta_Xdot",Xdot);
 74:     }
 75:   }
 76:   return(0);
 77: }

 79: static PetscErrorCode DMCoarsenHook_TSTheta(DM fine,DM coarse,void *ctx)
 80: {
 82:   return(0);
 83: }

 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: }

103: static PetscErrorCode DMSubDomainHook_TSTheta(DM dm,DM subdm,void *ctx)
104: {
106:   return(0);
107: }

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

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

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

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

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

130: static PetscErrorCode TSThetaEvaluateCostIntegral(TS ts)
131: {
132:   TS_Theta       *th = (TS_Theta*)ts->data;
133:   TS             quadts = ts->quadraturets;

137:   if (th->endpoint) {
138:     /* Evolve ts->vec_costintegral to compute integrals */
139:     if (th->Theta!=1.0) {
140:       TSComputeRHSFunction(quadts,th->ptime,th->X0,ts->vec_costintegrand);
141:       VecAXPY(quadts->vec_sol,th->time_step*(1.0-th->Theta),ts->vec_costintegrand);
142:     }
143:     TSComputeRHSFunction(quadts,ts->ptime,ts->vec_sol,ts->vec_costintegrand);
144:     VecAXPY(quadts->vec_sol,th->time_step*th->Theta,ts->vec_costintegrand);
145:   } else {
146:     TSComputeRHSFunction(quadts,th->stage_time,th->X,ts->vec_costintegrand);
147:     VecAXPY(quadts->vec_sol,th->time_step,ts->vec_costintegrand);
148:   }
149:   return(0);
150: }

152: static PetscErrorCode TSForwardCostIntegral_Theta(TS ts)
153: {
154:   TS_Theta       *th = (TS_Theta*)ts->data;
155:   TS             quadts = ts->quadraturets;

159:   /* backup cost integral */
160:   VecCopy(quadts->vec_sol,th->VecCostIntegral0);
161:   TSThetaEvaluateCostIntegral(ts);
162:   return(0);
163: }

165: static PetscErrorCode TSAdjointCostIntegral_Theta(TS ts)
166: {

170:   TSThetaEvaluateCostIntegral(ts);
171:   return(0);
172: }

174: static PetscErrorCode TSTheta_SNESSolve(TS ts,Vec b,Vec x)
175: {
176:   PetscInt       nits,lits;

180:   SNESSolve(ts->snes,b,x);
181:   SNESGetIterationNumber(ts->snes,&nits);
182:   SNESGetLinearSolveIterations(ts->snes,&lits);
183:   ts->snes_its += nits; ts->ksp_its += lits;
184:   return(0);
185: }

187: static PetscErrorCode TSStep_Theta(TS ts)
188: {
189:   TS_Theta       *th = (TS_Theta*)ts->data;
190:   PetscInt       rejections = 0;
191:   PetscBool      stageok,accept = PETSC_TRUE;
192:   PetscReal      next_time_step = ts->time_step;

196:   if (!ts->steprollback) {
197:     if (th->vec_sol_prev) { VecCopy(th->X0,th->vec_sol_prev); }
198:     VecCopy(ts->vec_sol,th->X0);
199:   }

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

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

207:     VecCopy(th->X0,th->X);
208:     if (th->extrapolate && !ts->steprestart) {
209:       VecAXPY(th->X,1/shift,th->Xdot);
210:     }
211:     if (th->endpoint) { /* This formulation assumes linear time-independent mass matrix */
212:       if (!th->affine) {VecDuplicate(ts->vec_sol,&th->affine);}
213:       VecZeroEntries(th->Xdot);
214:       TSComputeIFunction(ts,ts->ptime,th->X0,th->Xdot,th->affine,PETSC_FALSE);
215:       VecScale(th->affine,(th->Theta-1)/th->Theta);
216:     } else if (th->affine) { /* Just in case th->endpoint is changed between calls to TSStep_Theta() */
217:       VecZeroEntries(th->affine);
218:     }
219:     TSPreStage(ts,th->stage_time);
220:     TSTheta_SNESSolve(ts,th->affine,th->X);
221:     TSPostStage(ts,th->stage_time,0,&th->X);
222:     TSAdaptCheckStage(ts->adapt,ts,th->stage_time,th->X,&stageok);
223:     if (!stageok) goto reject_step;

225:     th->status = TS_STEP_PENDING;
226:     if (th->endpoint) {
227:       VecCopy(th->X,ts->vec_sol);
228:     } else {
229:       VecAXPBYPCZ(th->Xdot,-shift,shift,0,th->X0,th->X);
230:       VecAXPY(ts->vec_sol,ts->time_step,th->Xdot);
231:     }
232:     TSAdaptChoose(ts->adapt,ts,ts->time_step,NULL,&next_time_step,&accept);
233:     th->status = accept ? TS_STEP_COMPLETE : TS_STEP_INCOMPLETE;
234:     if (!accept) {
235:       VecCopy(th->X0,ts->vec_sol);
236:       ts->time_step = next_time_step;
237:       goto reject_step;
238:     }

240:     if (ts->forward_solve || ts->costintegralfwd) { /* Save the info for the later use in cost integral evaluation */
241:       th->ptime     = ts->ptime;
242:       th->time_step = ts->time_step;
243:     }

245:     ts->ptime += ts->time_step;
246:     ts->time_step = next_time_step;
247:     break;

249:   reject_step:
250:     ts->reject++; accept = PETSC_FALSE;
251:     if (!ts->reason && ++rejections > ts->max_reject && ts->max_reject >= 0) {
252:       ts->reason = TS_DIVERGED_STEP_REJECTED;
253:       PetscInfo2(ts,"Step=%D, step rejections %D greater than current TS allowed, stopping solve\n",ts->steps,rejections);
254:     }
255:   }
256:   return(0);
257: }

259: static PetscErrorCode TSAdjointStepBEuler_Private(TS ts)
260: {
261:   TS_Theta       *th = (TS_Theta*)ts->data;
262:   TS             quadts = ts->quadraturets;
263:   Vec            *VecsDeltaLam = th->VecsDeltaLam,*VecsDeltaMu = th->VecsDeltaMu,*VecsSensiTemp = th->VecsSensiTemp;
264:   Vec            *VecsDeltaLam2 = th->VecsDeltaLam2,*VecsDeltaMu2 = th->VecsDeltaMu2,*VecsSensi2Temp = th->VecsSensi2Temp;
265:   PetscInt       nadj;
266:   Mat            J,Jpre,quadJ = NULL,quadJp = NULL;
267:   KSP            ksp;
268:   PetscReal      shift;
269:   PetscScalar    *xarr;

273:   th->status = TS_STEP_INCOMPLETE;
274:   SNESGetKSP(ts->snes,&ksp);
275:   TSGetIJacobian(ts,&J,&Jpre,NULL,NULL);
276:   if (quadts) {
277:     TSGetRHSJacobian(quadts,&quadJ,NULL,NULL,NULL);
278:     TSGetRHSJacobianP(quadts,&quadJp,NULL,NULL);
279:   }

281:   /* If endpoint=1, th->ptime and th->X0 will be used; if endpoint=0, th->stage_time and th->X will be used. */
282:   th->stage_time = ts->ptime; /* time_step is negative*/
283:   th->ptime      = ts->ptime + ts->time_step;
284:   th->time_step  = -ts->time_step;

286:   /* Build RHS for first-order adjoint lambda_{n+1}/h + r_u^T(n+1) */
287:   if (quadts) {
288:     TSComputeRHSJacobian(quadts,th->stage_time,th->X,quadJ,NULL);
289:   }

291:   for (nadj=0; nadj<ts->numcost; nadj++) {
292:     VecCopy(ts->vecs_sensi[nadj],VecsSensiTemp[nadj]);
293:     VecScale(VecsSensiTemp[nadj],1./th->time_step); /* lambda_{n+1}/h */
294:     if (quadJ) {
295:       MatDenseGetColumn(quadJ,nadj,&xarr);
296:       VecPlaceArray(ts->vec_drdu_col,xarr);
297:       VecAXPY(VecsSensiTemp[nadj],1.,ts->vec_drdu_col);
298:       VecResetArray(ts->vec_drdu_col);
299:       MatDenseRestoreColumn(quadJ,&xarr);
300:     }
301:   }

303:   /* Build LHS for first-order adjoint */
304:   shift = 1./th->time_step;
305:   TSComputeIJacobian(ts,th->stage_time,th->X,th->Xdot,shift,J,Jpre,PETSC_FALSE);
306:   KSPSetOperators(ksp,J,Jpre);

308:   /* Solve stage equation LHS*lambda_s = RHS for first-order adjoint */
309:   for (nadj=0; nadj<ts->numcost; nadj++) {
310:     KSPConvergedReason kspreason;
311:     KSPSolveTranspose(ksp,VecsSensiTemp[nadj],VecsDeltaLam[nadj]);
312:     KSPGetConvergedReason(ksp,&kspreason);
313:     if (kspreason < 0) {
314:       ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
315:       PetscInfo2(ts,"Step=%D, %Dth cost function, transposed linear solve fails, stopping adjoint solve\n",ts->steps,nadj);
316:     }
317:   }

319:   if (ts->vecs_sensi2) { /* U_{n+1} */
320:     /* Get w1 at t_{n+1} from TLM matrix */
321:     MatDenseGetColumn(ts->mat_sensip,0,&xarr);
322:     VecPlaceArray(ts->vec_sensip_col,xarr);
323:     /* lambda_s^T F_UU w_1 */
324:     TSComputeIHessianProductFunctionUU(ts,th->stage_time,th->X,VecsDeltaLam,ts->vec_sensip_col,ts->vecs_fuu);
325:     /* lambda_s^T F_UP w_2 */
326:     TSComputeIHessianProductFunctionUP(ts,th->stage_time,th->X,VecsDeltaLam,ts->vec_dir,ts->vecs_fup);
327:     for (nadj=0; nadj<ts->numcost; nadj++) { /* compute the residual */
328:       VecCopy(ts->vecs_sensi2[nadj],VecsSensi2Temp[nadj]);
329:       VecScale(VecsSensi2Temp[nadj],shift);
330:       VecAXPY(VecsSensi2Temp[nadj],-1.,ts->vecs_fuu[nadj]);
331:       if (ts->vecs_fup) {
332:         VecAXPY(VecsSensi2Temp[nadj],-1.,ts->vecs_fup[nadj]);
333:       }
334:     }
335:     /* Solve stage equation LHS X = RHS for second-order adjoint */
336:     for (nadj=0; nadj<ts->numcost; nadj++) {
337:       KSPConvergedReason kspreason;
338:       KSPSolveTranspose(ksp,VecsSensi2Temp[nadj],VecsDeltaLam2[nadj]);
339:       KSPGetConvergedReason(ksp,&kspreason);
340:       if (kspreason < 0) {
341:         ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
342:         PetscInfo2(ts,"Step=%D, %Dth cost function, transposed linear solve fails, stopping adjoint solve\n",ts->steps,nadj);
343:       }
344:     }
345:   }

347:   /* Update sensitivities, and evaluate integrals if there is any */
348:   shift = 0.0;
349:   TSComputeIJacobian(ts,th->stage_time,th->X,th->Xdot,shift,J,Jpre,PETSC_FALSE); /* get -f_U */
350:   MatScale(J,-1.);
351:   for (nadj=0; nadj<ts->numcost; nadj++) {
352:     /* Add f_U \lambda_s to the original RHS */
353:     MatMultTransposeAdd(J,VecsDeltaLam[nadj],VecsSensiTemp[nadj],VecsSensiTemp[nadj]);
354:     VecScale(VecsSensiTemp[nadj],th->time_step);
355:     VecCopy(VecsSensiTemp[nadj],ts->vecs_sensi[nadj]);
356:     if (ts->vecs_sensi2) {
357:       MatMultTransposeAdd(J,VecsDeltaLam2[nadj],VecsSensi2Temp[nadj],VecsSensi2Temp[nadj]);
358:       VecScale(VecsSensi2Temp[nadj],th->time_step);
359:       VecCopy(VecsSensi2Temp[nadj],ts->vecs_sensi2[nadj]);
360:     }
361:   }
362:   if (ts->vecs_sensip) {
363:     TSComputeIJacobianP(ts,th->stage_time,th->X,th->Xdot,shift,ts->Jacp,PETSC_FALSE); /* get -f_p */
364:     if (quadts) {
365:       TSComputeRHSJacobianP(quadts,th->stage_time,th->X,quadJp);
366:     }
367:     if (ts->vecs_sensi2p) {
368:       /* lambda_s^T F_PU w_1 */
369:       TSComputeIHessianProductFunctionPU(ts,th->stage_time,th->X,VecsDeltaLam,ts->vec_sensip_col,ts->vecs_fpu);
370:       /* lambda_s^T F_PP w_2 */
371:       TSComputeIHessianProductFunctionPP(ts,th->stage_time,th->X,VecsDeltaLam,ts->vec_dir,ts->vecs_fpp);
372:     }

374:     for (nadj=0; nadj<ts->numcost; nadj++) {
375:       MatMultTranspose(ts->Jacp,VecsDeltaLam[nadj],VecsDeltaMu[nadj]);
376:       VecAXPY(ts->vecs_sensip[nadj],-th->time_step,VecsDeltaMu[nadj]);
377:       if (quadJp) {
378:         MatDenseGetColumn(quadJp,nadj,&xarr);
379:         VecPlaceArray(ts->vec_drdp_col,xarr);
380:         VecAXPY(ts->vecs_sensip[nadj],th->time_step,ts->vec_drdp_col);
381:         VecResetArray(ts->vec_drdp_col);
382:         MatDenseRestoreColumn(quadJp,&xarr);
383:       }
384:       if (ts->vecs_sensi2p) {
385:         MatMultTranspose(ts->Jacp,VecsDeltaLam2[nadj],VecsDeltaMu2[nadj]);
386:         VecAXPY(ts->vecs_sensi2p[nadj],-th->time_step,VecsDeltaMu2[nadj]);
387:         if (ts->vecs_fpu) {
388:           VecAXPY(ts->vecs_sensi2p[nadj],-th->time_step,ts->vecs_fpu[nadj]);
389:         }
390:         if (ts->vecs_fpp) {
391:           VecAXPY(ts->vecs_sensi2p[nadj],-th->time_step,ts->vecs_fpp[nadj]);
392:         }
393:       }
394:     }
395:   }

397:   if (ts->vecs_sensi2) {
398:     VecResetArray(ts->vec_sensip_col);
399:     MatDenseRestoreColumn(ts->mat_sensip,&xarr);
400:   }
401:   th->status = TS_STEP_COMPLETE;
402:   return(0);
403: }

405: static PetscErrorCode TSAdjointStep_Theta(TS ts)
406: {
407:   TS_Theta       *th = (TS_Theta*)ts->data;
408:   TS             quadts = ts->quadraturets;
409:   Vec            *VecsDeltaLam = th->VecsDeltaLam,*VecsDeltaMu = th->VecsDeltaMu,*VecsSensiTemp = th->VecsSensiTemp;
410:   Vec            *VecsDeltaLam2 = th->VecsDeltaLam2,*VecsDeltaMu2 = th->VecsDeltaMu2,*VecsSensi2Temp = th->VecsSensi2Temp;
411:   PetscInt       nadj;
412:   Mat            J,Jpre,quadJ = NULL,quadJp = NULL;
413:   KSP            ksp;
414:   PetscReal      shift;
415:   PetscScalar    *xarr;

419:   if (th->Theta == 1.) {
420:     TSAdjointStepBEuler_Private(ts);
421:     return(0);
422:   }
423:   th->status = TS_STEP_INCOMPLETE;
424:   SNESGetKSP(ts->snes,&ksp);
425:   TSGetIJacobian(ts,&J,&Jpre,NULL,NULL);
426:   if (quadts) {
427:     TSGetRHSJacobian(quadts,&quadJ,NULL,NULL,NULL);
428:     TSGetRHSJacobianP(quadts,&quadJp,NULL,NULL);
429:   }
430:   /* If endpoint=1, th->ptime and th->X0 will be used; if endpoint=0, th->stage_time and th->X will be used. */
431:   th->stage_time = th->endpoint ? ts->ptime : (ts->ptime+(1.-th->Theta)*ts->time_step); /* time_step is negative*/
432:   th->ptime      = ts->ptime + ts->time_step;
433:   th->time_step  = -ts->time_step;

435:   /* Build RHS for first-order adjoint */
436:   /* Cost function has an integral term */
437:   if (quadts) {
438:     if (th->endpoint) {
439:       TSComputeRHSJacobian(quadts,th->stage_time,ts->vec_sol,quadJ,NULL);
440:     } else {
441:       TSComputeRHSJacobian(quadts,th->stage_time,th->X,quadJ,NULL);
442:     }
443:   }

445:   for (nadj=0; nadj<ts->numcost; nadj++) {
446:     VecCopy(ts->vecs_sensi[nadj],VecsSensiTemp[nadj]);
447:     VecScale(VecsSensiTemp[nadj],1./(th->Theta*th->time_step));
448:     if (quadJ) {
449:       MatDenseGetColumn(quadJ,nadj,&xarr);
450:       VecPlaceArray(ts->vec_drdu_col,xarr);
451:       VecAXPY(VecsSensiTemp[nadj],1.,ts->vec_drdu_col);
452:       VecResetArray(ts->vec_drdu_col);
453:       MatDenseRestoreColumn(quadJ,&xarr);
454:     }
455:   }

457:   /* Build LHS for first-order adjoint */
458:   shift = 1./(th->Theta*th->time_step);
459:   if (th->endpoint) {
460:     TSComputeIJacobian(ts,th->stage_time,ts->vec_sol,th->Xdot,shift,J,Jpre,PETSC_FALSE);
461:   } else {
462:     TSComputeIJacobian(ts,th->stage_time,th->X,th->Xdot,shift,J,Jpre,PETSC_FALSE);
463:   }
464:   KSPSetOperators(ksp,J,Jpre);

466:   /* Solve stage equation LHS*lambda_s = RHS for first-order adjoint */
467:   for (nadj=0; nadj<ts->numcost; nadj++) {
468:     KSPConvergedReason kspreason;
469:     KSPSolveTranspose(ksp,VecsSensiTemp[nadj],VecsDeltaLam[nadj]);
470:     KSPGetConvergedReason(ksp,&kspreason);
471:     if (kspreason < 0) {
472:       ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
473:       PetscInfo2(ts,"Step=%D, %Dth cost function, transposed linear solve fails, stopping adjoint solve\n",ts->steps,nadj);
474:     }
475:   }

477:   /* Second-order adjoint */
478:   if (ts->vecs_sensi2) { /* U_{n+1} */
479:     if (!th->endpoint) SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_SUP,"Operation not implemented in TS_Theta");
480:     /* Get w1 at t_{n+1} from TLM matrix */
481:     MatDenseGetColumn(ts->mat_sensip,0,&xarr);
482:     VecPlaceArray(ts->vec_sensip_col,xarr);
483:     /* lambda_s^T F_UU w_1 */
484:     TSComputeIHessianProductFunctionUU(ts,th->stage_time,ts->vec_sol,VecsDeltaLam,ts->vec_sensip_col,ts->vecs_fuu);
485:     VecResetArray(ts->vec_sensip_col);
486:     MatDenseRestoreColumn(ts->mat_sensip,&xarr);
487:     /* lambda_s^T F_UP w_2 */
488:     TSComputeIHessianProductFunctionUP(ts,th->stage_time,ts->vec_sol,VecsDeltaLam,ts->vec_dir,ts->vecs_fup);
489:     for (nadj=0; nadj<ts->numcost; nadj++) { /* compute the residual */
490:       VecCopy(ts->vecs_sensi2[nadj],VecsSensi2Temp[nadj]);
491:       VecScale(VecsSensi2Temp[nadj],shift);
492:       VecAXPY(VecsSensi2Temp[nadj],-1.,ts->vecs_fuu[nadj]);
493:       if (ts->vecs_fup) {
494:         VecAXPY(VecsSensi2Temp[nadj],-1.,ts->vecs_fup[nadj]);
495:       }
496:     }
497:     /* Solve stage equation LHS X = RHS for second-order adjoint */
498:     for (nadj=0; nadj<ts->numcost; nadj++) {
499:       KSPConvergedReason kspreason;
500:       KSPSolveTranspose(ksp,VecsSensi2Temp[nadj],VecsDeltaLam2[nadj]);
501:       KSPGetConvergedReason(ksp,&kspreason);
502:       if (kspreason < 0) {
503:         ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
504:         PetscInfo2(ts,"Step=%D, %Dth cost function, transposed linear solve fails, stopping adjoint solve\n",ts->steps,nadj);
505:       }
506:     }
507:   }

509:   /* Update sensitivities, and evaluate integrals if there is any */
510:   if(th->endpoint) { /* two-stage Theta methods with th->Theta!=1, th->Theta==1 leads to BEuler */
511:     shift = 1./((th->Theta-1.)*th->time_step);
512:     TSComputeIJacobian(ts,th->ptime,th->X0,th->Xdot,shift,J,Jpre,PETSC_FALSE);
513:     /* R_U at t_n */
514:     if (quadts) {
515:       TSComputeRHSJacobian(quadts,th->ptime,th->X0,quadJ,NULL);
516:     }
517:     for (nadj=0; nadj<ts->numcost; nadj++) {
518:       MatMultTranspose(J,VecsDeltaLam[nadj],ts->vecs_sensi[nadj]);
519:       if (quadJ) {
520:         MatDenseGetColumn(quadJ,nadj,&xarr);
521:         VecPlaceArray(ts->vec_drdu_col,xarr);
522:         VecAXPY(ts->vecs_sensi[nadj],-1.,ts->vec_drdu_col);
523:         VecResetArray(ts->vec_drdu_col);
524:         MatDenseRestoreColumn(quadJ,&xarr);
525:       }
526:       VecScale(ts->vecs_sensi[nadj],1./shift);
527:     }

529:     /* Second-order adjoint */
530:     if (ts->vecs_sensi2) { /* U_n */
531:       /* Get w1 at t_n from TLM matrix */
532:       MatDenseGetColumn(th->MatFwdSensip0,0,&xarr);
533:       VecPlaceArray(ts->vec_sensip_col,xarr);
534:       /* lambda_s^T F_UU w_1 */
535:       TSComputeIHessianProductFunctionUU(ts,th->ptime,th->X0,VecsDeltaLam,ts->vec_sensip_col,ts->vecs_fuu);
536:       VecResetArray(ts->vec_sensip_col);
537:       MatDenseRestoreColumn(th->MatFwdSensip0,&xarr);
538:       /* lambda_s^T F_UU w_2 */
539:       TSComputeIHessianProductFunctionUP(ts,th->ptime,th->X0,VecsDeltaLam,ts->vec_dir,ts->vecs_fup);
540:       for (nadj=0; nadj<ts->numcost; nadj++) {
541:         /* M^T Lambda_s + h(1-theta) F_U^T Lambda_s + h(1-theta) lambda_s^T F_UU w_1 + lambda_s^T F_UP w_2  */
542:         MatMultTranspose(J,VecsDeltaLam2[nadj],ts->vecs_sensi2[nadj]);
543:         VecAXPY(ts->vecs_sensi2[nadj],1.,ts->vecs_fuu[nadj]);
544:         if (ts->vecs_fup) {
545:           VecAXPY(ts->vecs_sensi2[nadj],1.,ts->vecs_fup[nadj]);
546:         }
547:         VecScale(ts->vecs_sensi2[nadj],1./shift);
548:       }
549:     }

551:     if (ts->vecs_sensip) { /* sensitivities wrt parameters */
552:       /* U_{n+1} */
553:       shift = -1./(th->Theta*th->time_step);
554:       TSComputeIJacobianP(ts,th->stage_time,ts->vec_sol,th->Xdot,shift,ts->Jacp,PETSC_FALSE);
555:       if (quadts) {
556:         TSComputeRHSJacobianP(quadts,th->stage_time,ts->vec_sol,quadJp);
557:       }
558:       for (nadj=0; nadj<ts->numcost; nadj++) {
559:         MatMultTranspose(ts->Jacp,VecsDeltaLam[nadj],VecsDeltaMu[nadj]);
560:         VecAXPY(ts->vecs_sensip[nadj],-th->time_step*th->Theta,VecsDeltaMu[nadj]);
561:       }
562:       if (ts->vecs_sensi2p) { /* second-order */
563:         /* Get w1 at t_{n+1} from TLM matrix */
564:         MatDenseGetColumn(ts->mat_sensip,0,&xarr);
565:         VecPlaceArray(ts->vec_sensip_col,xarr);
566:         /* lambda_s^T F_PU w_1 */
567:         TSComputeIHessianProductFunctionPU(ts,th->stage_time,ts->vec_sol,VecsDeltaLam,ts->vec_sensip_col,ts->vecs_fpu);
568:         VecResetArray(ts->vec_sensip_col);
569:         MatDenseRestoreColumn(ts->mat_sensip,&xarr);

571:         /* lambda_s^T F_PP w_2 */
572:         TSComputeIHessianProductFunctionPP(ts,th->stage_time,ts->vec_sol,VecsDeltaLam,ts->vec_dir,ts->vecs_fpp);
573:         for (nadj=0; nadj<ts->numcost; nadj++) {
574:           /* Mu2 <- Mu2 + h theta F_P^T Lambda_s + h theta (lambda_s^T F_UU w_1 + lambda_s^T F_UP w_2)  */
575:           MatMultTranspose(ts->Jacp,VecsDeltaLam2[nadj],VecsDeltaMu2[nadj]);
576:           VecAXPY(ts->vecs_sensi2p[nadj],-th->time_step*th->Theta,VecsDeltaMu2[nadj]);
577:           if (ts->vecs_fpu) {
578:             VecAXPY(ts->vecs_sensi2p[nadj],-th->time_step*th->Theta,ts->vecs_fpu[nadj]);
579:           }
580:           if (ts->vecs_fpp) {
581:             VecAXPY(ts->vecs_sensi2p[nadj],-th->time_step*th->Theta,ts->vecs_fpp[nadj]);
582:           }
583:         }
584:       }

586:       /* U_s */
587:       shift = 1./((th->Theta-1.0)*th->time_step);
588:       TSComputeIJacobianP(ts,th->ptime,th->X0,th->Xdot,shift,ts->Jacp,PETSC_FALSE);
589:       if (quadts) {
590:         TSComputeRHSJacobianP(quadts,th->ptime,th->X0,quadJp);
591:       }
592:       for (nadj=0; nadj<ts->numcost; nadj++) {
593:         MatMultTranspose(ts->Jacp,VecsDeltaLam[nadj],VecsDeltaMu[nadj]);
594:         VecAXPY(ts->vecs_sensip[nadj],-th->time_step*(1.0-th->Theta),VecsDeltaMu[nadj]);
595:         if (ts->vecs_sensi2p) { /* second-order */
596:           /* Get w1 at t_n from TLM matrix */
597:           MatDenseGetColumn(th->MatFwdSensip0,0,&xarr);
598:           VecPlaceArray(ts->vec_sensip_col,xarr);
599:           /* lambda_s^T F_PU w_1 */
600:           TSComputeIHessianProductFunctionPU(ts,th->ptime,th->X0,VecsDeltaLam,ts->vec_sensip_col,ts->vecs_fpu);
601:           VecResetArray(ts->vec_sensip_col);
602:           MatDenseRestoreColumn(th->MatFwdSensip0,&xarr);
603:           /* lambda_s^T F_PP w_2 */
604:           TSComputeIHessianProductFunctionPP(ts,th->ptime,th->X0,VecsDeltaLam,ts->vec_dir,ts->vecs_fpp);
605:           for (nadj=0; nadj<ts->numcost; nadj++) {
606:             /* Mu2 <- Mu2 + h(1-theta) F_P^T Lambda_s + h(1-theta) (lambda_s^T F_UU w_1 + lambda_s^T F_UP w_2) */
607:             MatMultTranspose(ts->Jacp,VecsDeltaLam2[nadj],VecsDeltaMu2[nadj]);
608:             VecAXPY(ts->vecs_sensi2p[nadj],-th->time_step*(1.0-th->Theta),VecsDeltaMu2[nadj]);
609:             if (ts->vecs_fpu) {
610:               VecAXPY(ts->vecs_sensi2p[nadj],-th->time_step*(1.0-th->Theta),ts->vecs_fpu[nadj]);
611:             }
612:             if (ts->vecs_fpp) {
613:               VecAXPY(ts->vecs_sensi2p[nadj],-th->time_step*(1.0-th->Theta),ts->vecs_fpp[nadj]);
614:             }
615:           }
616:         }
617:       }
618:     }
619:   } else { /* one-stage case */
620:     shift = 0.0;
621:     TSComputeIJacobian(ts,th->stage_time,th->X,th->Xdot,shift,J,Jpre,PETSC_FALSE); /* get -f_y */
622:     if (quadts) {
623:       TSComputeRHSJacobian(quadts,th->stage_time,th->X,quadJ,NULL);
624:     }
625:     for (nadj=0; nadj<ts->numcost; nadj++) {
626:       MatMultTranspose(J,VecsDeltaLam[nadj],VecsSensiTemp[nadj]);
627:       VecAXPY(ts->vecs_sensi[nadj],-th->time_step,VecsSensiTemp[nadj]);
628:       if (quadJ) {
629:         MatDenseGetColumn(quadJ,nadj,&xarr);
630:         VecPlaceArray(ts->vec_drdu_col,xarr);
631:         VecAXPY(ts->vecs_sensi[nadj],th->time_step,ts->vec_drdu_col);
632:         VecResetArray(ts->vec_drdu_col);
633:         MatDenseRestoreColumn(quadJ,&xarr);
634:       }
635:     }
636:     if (ts->vecs_sensip) {
637:       TSComputeIJacobianP(ts,th->stage_time,th->X,th->Xdot,shift,ts->Jacp,PETSC_FALSE);
638:       if (quadts) {
639:         TSComputeRHSJacobianP(quadts,th->stage_time,th->X,quadJp);
640:       }
641:       for (nadj=0; nadj<ts->numcost; nadj++) {
642:         MatMultTranspose(ts->Jacp,VecsDeltaLam[nadj],VecsDeltaMu[nadj]);
643:         VecAXPY(ts->vecs_sensip[nadj],-th->time_step,VecsDeltaMu[nadj]);
644:         if (quadJp) {
645:           MatDenseGetColumn(quadJp,nadj,&xarr);
646:           VecPlaceArray(ts->vec_drdp_col,xarr);
647:           VecAXPY(ts->vecs_sensip[nadj],th->time_step,ts->vec_drdp_col);
648:           VecResetArray(ts->vec_drdp_col);
649:           MatDenseRestoreColumn(quadJp,&xarr);
650:         }
651:       }
652:     }
653:   }

655:   th->status = TS_STEP_COMPLETE;
656:   return(0);
657: }

659: static PetscErrorCode TSInterpolate_Theta(TS ts,PetscReal t,Vec X)
660: {
661:   TS_Theta       *th = (TS_Theta*)ts->data;
662:   PetscReal      dt  = t - ts->ptime;

666:   VecCopy(ts->vec_sol,th->X);
667:   if (th->endpoint) dt *= th->Theta;
668:   VecWAXPY(X,dt,th->Xdot,th->X);
669:   return(0);
670: }

672: static PetscErrorCode TSEvaluateWLTE_Theta(TS ts,NormType wnormtype,PetscInt *order,PetscReal *wlte)
673: {
674:   TS_Theta       *th = (TS_Theta*)ts->data;
675:   Vec            X = ts->vec_sol;      /* X = solution */
676:   Vec            Y = th->vec_lte_work; /* Y = X + LTE  */
677:   PetscReal      wltea,wlter;

681:   if (!th->vec_sol_prev) {*wlte = -1; return(0);}
682:   /* Cannot compute LTE in first step or in restart after event */
683:   if (ts->steprestart) {*wlte = -1; return(0);}
684:   /* Compute LTE using backward differences with non-constant time step */
685:   {
686:     PetscReal   h = ts->time_step, h_prev = ts->ptime - ts->ptime_prev;
687:     PetscReal   a = 1 + h_prev/h;
688:     PetscScalar scal[3]; Vec vecs[3];
689:     scal[0] = +1/a; scal[1] = -1/(a-1); scal[2] = +1/(a*(a-1));
690:     vecs[0] = X;    vecs[1] = th->X0;   vecs[2] = th->vec_sol_prev;
691:     VecCopy(X,Y);
692:     VecMAXPY(Y,3,scal,vecs);
693:     TSErrorWeightedNorm(ts,X,Y,wnormtype,wlte,&wltea,&wlter);
694:   }
695:   if (order) *order = 2;
696:   return(0);
697: }

699: static PetscErrorCode TSRollBack_Theta(TS ts)
700: {
701:   TS_Theta       *th = (TS_Theta*)ts->data;
702:   TS             quadts = ts->quadraturets;

706:   VecCopy(th->X0,ts->vec_sol);
707:   if (quadts && ts->costintegralfwd) {
708:     VecCopy(th->VecCostIntegral0,quadts->vec_sol);
709:   }
710:   th->status = TS_STEP_INCOMPLETE;
711:   if (ts->mat_sensip) {
712:     MatCopy(th->MatFwdSensip0,ts->mat_sensip,SAME_NONZERO_PATTERN);
713:   }
714:   if (quadts && quadts->mat_sensip) {
715:     MatCopy(th->MatIntegralSensip0,quadts->mat_sensip,SAME_NONZERO_PATTERN);
716:   }
717:   return(0);
718: }

720: static PetscErrorCode TSForwardStep_Theta(TS ts)
721: {
722:   TS_Theta       *th = (TS_Theta*)ts->data;
723:   TS             quadts = ts->quadraturets;
724:   Mat            MatDeltaFwdSensip = th->MatDeltaFwdSensip;
725:   Vec            VecDeltaFwdSensipCol = th->VecDeltaFwdSensipCol;
726:   PetscInt       ntlm;
727:   KSP            ksp;
728:   Mat            J,Jpre,quadJ = NULL,quadJp = NULL;
729:   PetscReal      shift;
730:   PetscScalar    *barr,*xarr;

734:   MatCopy(ts->mat_sensip,th->MatFwdSensip0,SAME_NONZERO_PATTERN);

736:   if (quadts && quadts->mat_sensip) {
737:     MatCopy(quadts->mat_sensip,th->MatIntegralSensip0,SAME_NONZERO_PATTERN);
738:   }
739:   SNESGetKSP(ts->snes,&ksp);
740:   TSGetIJacobian(ts,&J,&Jpre,NULL,NULL);
741:   if (quadts) {
742:     TSGetRHSJacobian(quadts,&quadJ,NULL,NULL,NULL);
743:     TSGetRHSJacobianP(quadts,&quadJp,NULL,NULL);
744:   }

746:   /* Build RHS */
747:   if (th->endpoint) { /* 2-stage method*/
748:     shift = 1./((th->Theta-1.)*th->time_step);
749:     TSComputeIJacobian(ts,th->ptime,th->X0,th->Xdot,shift,J,Jpre,PETSC_FALSE);
750:     MatMatMult(J,ts->mat_sensip,MAT_REUSE_MATRIX,PETSC_DEFAULT,&MatDeltaFwdSensip);
751:     MatScale(MatDeltaFwdSensip,(th->Theta-1.)/th->Theta);

753:     /* Add the f_p forcing terms */
754:     if (ts->Jacp) {
755:       TSComputeIJacobianP(ts,th->ptime,th->X0,th->Xdot,shift,ts->Jacp,PETSC_FALSE);
756:       MatAXPY(MatDeltaFwdSensip,(th->Theta-1.)/th->Theta,ts->Jacp,SUBSET_NONZERO_PATTERN);
757:       TSComputeIJacobianP(ts,th->stage_time,ts->vec_sol,th->Xdot,shift,ts->Jacp,PETSC_FALSE);
758:       MatAXPY(MatDeltaFwdSensip,-1.,ts->Jacp,SUBSET_NONZERO_PATTERN);
759:     }
760:   } else { /* 1-stage method */
761:     shift = 0.0;
762:     TSComputeIJacobian(ts,th->stage_time,th->X,th->Xdot,shift,J,Jpre,PETSC_FALSE);
763:     MatMatMult(J,ts->mat_sensip,MAT_REUSE_MATRIX,PETSC_DEFAULT,&MatDeltaFwdSensip);
764:     MatScale(MatDeltaFwdSensip,-1.);

766:     /* Add the f_p forcing terms */
767:     if (ts->Jacp) {
768:       TSComputeIJacobianP(ts,th->stage_time,th->X,th->Xdot,shift,ts->Jacp,PETSC_FALSE);
769:       MatAXPY(MatDeltaFwdSensip,-1.,ts->Jacp,SUBSET_NONZERO_PATTERN);
770:     }
771:   }

773:   /* Build LHS */
774:   shift = 1/(th->Theta*th->time_step);
775:   if (th->endpoint) {
776:     TSComputeIJacobian(ts,th->stage_time,ts->vec_sol,th->Xdot,shift,J,Jpre,PETSC_FALSE);
777:   } else {
778:     TSComputeIJacobian(ts,th->stage_time,th->X,th->Xdot,shift,J,Jpre,PETSC_FALSE);
779:   }
780:   KSPSetOperators(ksp,J,Jpre);

782:   /*
783:     Evaluate the first stage of integral gradients with the 2-stage method:
784:     drdu|t_n*S(t_n) + drdp|t_n
785:     This is done before the linear solve because the sensitivity variable S(t_n) will be propagated to S(t_{n+1})
786:   */
787:   if (th->endpoint) { /* 2-stage method only */
788:     if (quadts && quadts->mat_sensip) {
789:       TSComputeRHSJacobian(quadts,th->ptime,th->X0,quadJ,NULL);
790:       TSComputeRHSJacobianP(quadts,th->ptime,th->X0,quadJp);
791:       MatTransposeMatMult(ts->mat_sensip,quadJ,MAT_REUSE_MATRIX,PETSC_DEFAULT,&th->MatIntegralSensipTemp);
792:       MatAXPY(th->MatIntegralSensipTemp,1,quadJp,SAME_NONZERO_PATTERN);
793:       MatAXPY(quadts->mat_sensip,th->time_step*(1.-th->Theta),th->MatIntegralSensipTemp,SAME_NONZERO_PATTERN);
794:     }
795:   }

797:   /* Solve the tangent linear equation for forward sensitivities to parameters */
798:   for (ntlm=0; ntlm<th->num_tlm; ntlm++) {
799:     KSPConvergedReason kspreason;
800:     MatDenseGetColumn(MatDeltaFwdSensip,ntlm,&barr);
801:     VecPlaceArray(VecDeltaFwdSensipCol,barr);
802:     if (th->endpoint) {
803:       MatDenseGetColumn(ts->mat_sensip,ntlm,&xarr);
804:       VecPlaceArray(ts->vec_sensip_col,xarr);
805:       KSPSolve(ksp,VecDeltaFwdSensipCol,ts->vec_sensip_col);
806:       VecResetArray(ts->vec_sensip_col);
807:       MatDenseRestoreColumn(ts->mat_sensip,&xarr);
808:     } else {
809:       KSPSolve(ksp,VecDeltaFwdSensipCol,VecDeltaFwdSensipCol);
810:     }
811:     KSPGetConvergedReason(ksp,&kspreason);
812:     if (kspreason < 0) {
813:       ts->reason = TSFORWARD_DIVERGED_LINEAR_SOLVE;
814:       PetscInfo2(ts,"Step=%D, %Dth tangent linear solve, linear solve fails, stopping tangent linear solve\n",ts->steps,ntlm);
815:     }
816:     VecResetArray(VecDeltaFwdSensipCol);
817:     MatDenseRestoreColumn(MatDeltaFwdSensip,&barr);
818:   }


821:   /*
822:     Evaluate the second stage of integral gradients with the 2-stage method:
823:     drdu|t_{n+1}*S(t_{n+1}) + drdp|t_{n+1}
824:   */
825:   if (quadts && quadts->mat_sensip) {
826:     if (!th->endpoint) {
827:       MatAXPY(ts->mat_sensip,1,MatDeltaFwdSensip,SAME_NONZERO_PATTERN); /* stage sensitivity */
828:       TSComputeRHSJacobian(quadts,th->stage_time,th->X,quadJ,NULL);
829:       TSComputeRHSJacobianP(quadts,th->stage_time,th->X,quadJp);
830:       MatTransposeMatMult(ts->mat_sensip,quadJ,MAT_REUSE_MATRIX,PETSC_DEFAULT,&th->MatIntegralSensipTemp);
831:       MatAXPY(th->MatIntegralSensipTemp,1,quadJp,SAME_NONZERO_PATTERN);
832:       MatAXPY(quadts->mat_sensip,th->time_step,th->MatIntegralSensipTemp,SAME_NONZERO_PATTERN);
833:       MatAXPY(ts->mat_sensip,(1.-th->Theta)/th->Theta,MatDeltaFwdSensip,SAME_NONZERO_PATTERN);
834:     } else {
835:       TSComputeRHSJacobian(quadts,th->stage_time,ts->vec_sol,quadJ,NULL);
836:       TSComputeRHSJacobianP(quadts,th->stage_time,ts->vec_sol,quadJp);
837:       MatTransposeMatMult(ts->mat_sensip,quadJ,MAT_REUSE_MATRIX,PETSC_DEFAULT,&th->MatIntegralSensipTemp);
838:       MatAXPY(th->MatIntegralSensipTemp,1,quadJp,SAME_NONZERO_PATTERN);
839:       MatAXPY(quadts->mat_sensip,th->time_step*th->Theta,th->MatIntegralSensipTemp,SAME_NONZERO_PATTERN);
840:     }
841:   } else {
842:     if (!th->endpoint) {
843:       MatAXPY(ts->mat_sensip,1./th->Theta,MatDeltaFwdSensip,SAME_NONZERO_PATTERN);
844:     }
845:   }
846:   return(0);
847: }

849: static PetscErrorCode TSForwardGetStages_Theta(TS ts,PetscInt *ns,Mat **stagesensip)
850: {
851:   TS_Theta *th = (TS_Theta*)ts->data;

854:   if (ns) *ns = 1;
855:   if (stagesensip) *stagesensip = th->endpoint ? &(th->MatFwdSensip0) : &(th->MatDeltaFwdSensip);
856:   return(0);
857: }

859: /*------------------------------------------------------------*/
860: static PetscErrorCode TSReset_Theta(TS ts)
861: {
862:   TS_Theta       *th = (TS_Theta*)ts->data;

866:   VecDestroy(&th->X);
867:   VecDestroy(&th->Xdot);
868:   VecDestroy(&th->X0);
869:   VecDestroy(&th->affine);

871:   VecDestroy(&th->vec_sol_prev);
872:   VecDestroy(&th->vec_lte_work);

874:   VecDestroy(&th->VecCostIntegral0);
875:   return(0);
876: }

878: static PetscErrorCode TSAdjointReset_Theta(TS ts)
879: {
880:   TS_Theta       *th = (TS_Theta*)ts->data;

884:   VecDestroyVecs(ts->numcost,&th->VecsDeltaLam);
885:   VecDestroyVecs(ts->numcost,&th->VecsDeltaMu);
886:   VecDestroyVecs(ts->numcost,&th->VecsDeltaLam2);
887:   VecDestroyVecs(ts->numcost,&th->VecsDeltaMu2);
888:   VecDestroyVecs(ts->numcost,&th->VecsSensiTemp);
889:   VecDestroyVecs(ts->numcost,&th->VecsSensi2Temp);
890:   return(0);
891: }

893: static PetscErrorCode TSDestroy_Theta(TS ts)
894: {

898:   TSReset_Theta(ts);
899:   if (ts->dm) {
900:     DMCoarsenHookRemove(ts->dm,DMCoarsenHook_TSTheta,DMRestrictHook_TSTheta,ts);
901:     DMSubDomainHookRemove(ts->dm,DMSubDomainHook_TSTheta,DMSubDomainRestrictHook_TSTheta,ts);
902:   }
903:   PetscFree(ts->data);
904:   PetscObjectComposeFunction((PetscObject)ts,"TSThetaGetTheta_C",NULL);
905:   PetscObjectComposeFunction((PetscObject)ts,"TSThetaSetTheta_C",NULL);
906:   PetscObjectComposeFunction((PetscObject)ts,"TSThetaGetEndpoint_C",NULL);
907:   PetscObjectComposeFunction((PetscObject)ts,"TSThetaSetEndpoint_C",NULL);
908:   return(0);
909: }

911: /*
912:   This defines the nonlinear equation that is to be solved with SNES
913:   G(U) = F[t0+Theta*dt, U, (U-U0)*shift] = 0
914: */
915: static PetscErrorCode SNESTSFormFunction_Theta(SNES snes,Vec x,Vec y,TS ts)
916: {
917:   TS_Theta       *th = (TS_Theta*)ts->data;
919:   Vec            X0,Xdot;
920:   DM             dm,dmsave;
921:   PetscReal      shift = 1/(th->Theta*ts->time_step);

924:   SNESGetDM(snes,&dm);
925:   /* When using the endpoint variant, this is actually 1/Theta * Xdot */
926:   TSThetaGetX0AndXdot(ts,dm,&X0,&Xdot);
927:   VecAXPBYPCZ(Xdot,-shift,shift,0,X0,x);

929:   /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */
930:   dmsave = ts->dm;
931:   ts->dm = dm;
932:   TSComputeIFunction(ts,th->stage_time,x,Xdot,y,PETSC_FALSE);
933:   ts->dm = dmsave;
934:   TSThetaRestoreX0AndXdot(ts,dm,&X0,&Xdot);
935:   return(0);
936: }

938: static PetscErrorCode SNESTSFormJacobian_Theta(SNES snes,Vec x,Mat A,Mat B,TS ts)
939: {
940:   TS_Theta       *th = (TS_Theta*)ts->data;
942:   Vec            Xdot;
943:   DM             dm,dmsave;
944:   PetscReal      shift = 1/(th->Theta*ts->time_step);

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

951:   dmsave = ts->dm;
952:   ts->dm = dm;
953:   TSComputeIJacobian(ts,th->stage_time,x,Xdot,shift,A,B,PETSC_FALSE);
954:   ts->dm = dmsave;
955:   TSThetaRestoreX0AndXdot(ts,dm,NULL,&Xdot);
956:   return(0);
957: }

959: static PetscErrorCode TSForwardSetUp_Theta(TS ts)
960: {
961:   TS_Theta       *th = (TS_Theta*)ts->data;
962:   TS             quadts = ts->quadraturets;

966:   /* combine sensitivities to parameters and sensitivities to initial values into one array */
967:   th->num_tlm = ts->num_parameters;
968:   MatDuplicate(ts->mat_sensip,MAT_DO_NOT_COPY_VALUES,&th->MatDeltaFwdSensip);
969:   if (quadts && quadts->mat_sensip) {
970:     MatDuplicate(quadts->mat_sensip,MAT_DO_NOT_COPY_VALUES,&th->MatIntegralSensipTemp);
971:     MatDuplicate(quadts->mat_sensip,MAT_DO_NOT_COPY_VALUES,&th->MatIntegralSensip0);
972:   }
973:   /* backup sensitivity results for roll-backs */
974:   MatDuplicate(ts->mat_sensip,MAT_DO_NOT_COPY_VALUES,&th->MatFwdSensip0);

976:   VecDuplicate(ts->vec_sol,&th->VecDeltaFwdSensipCol);
977:   return(0);
978: }

980: static PetscErrorCode TSForwardReset_Theta(TS ts)
981: {
982:   TS_Theta       *th = (TS_Theta*)ts->data;
983:   TS             quadts = ts->quadraturets;

987:   if (quadts && quadts->mat_sensip) {
988:     MatDestroy(&th->MatIntegralSensipTemp);
989:     MatDestroy(&th->MatIntegralSensip0);
990:   }
991:   VecDestroy(&th->VecDeltaFwdSensipCol);
992:   MatDestroy(&th->MatDeltaFwdSensip);
993:   MatDestroy(&th->MatFwdSensip0);
994:   return(0);
995: }

997: static PetscErrorCode TSSetUp_Theta(TS ts)
998: {
999:   TS_Theta       *th = (TS_Theta*)ts->data;
1000:   TS             quadts = ts->quadraturets;
1001:   PetscBool      match;

1005:   if (!th->VecCostIntegral0 && quadts && ts->costintegralfwd) { /* back up cost integral */
1006:     VecDuplicate(quadts->vec_sol,&th->VecCostIntegral0);
1007:   }
1008:   if (!th->X) {
1009:     VecDuplicate(ts->vec_sol,&th->X);
1010:   }
1011:   if (!th->Xdot) {
1012:     VecDuplicate(ts->vec_sol,&th->Xdot);
1013:   }
1014:   if (!th->X0) {
1015:     VecDuplicate(ts->vec_sol,&th->X0);
1016:   }
1017:   if (th->endpoint) {
1018:     VecDuplicate(ts->vec_sol,&th->affine);
1019:   }

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

1023:   TSGetDM(ts,&ts->dm);
1024:   DMCoarsenHookAdd(ts->dm,DMCoarsenHook_TSTheta,DMRestrictHook_TSTheta,ts);
1025:   DMSubDomainHookAdd(ts->dm,DMSubDomainHook_TSTheta,DMSubDomainRestrictHook_TSTheta,ts);

1027:   TSGetAdapt(ts,&ts->adapt);
1028:   TSAdaptCandidatesClear(ts->adapt);
1029:   PetscObjectTypeCompare((PetscObject)ts->adapt,TSADAPTNONE,&match);
1030:   if (!match) {
1031:     VecDuplicate(ts->vec_sol,&th->vec_sol_prev);
1032:     VecDuplicate(ts->vec_sol,&th->vec_lte_work);
1033:   }
1034:   TSGetSNES(ts,&ts->snes);
1035:   return(0);
1036: }

1038: /*------------------------------------------------------------*/

1040: static PetscErrorCode TSAdjointSetUp_Theta(TS ts)
1041: {
1042:   TS_Theta       *th = (TS_Theta*)ts->data;

1046:   VecDuplicateVecs(ts->vecs_sensi[0],ts->numcost,&th->VecsDeltaLam);
1047:   VecDuplicateVecs(ts->vecs_sensi[0],ts->numcost,&th->VecsSensiTemp);
1048:   if (ts->vecs_sensip) {
1049:     VecDuplicateVecs(ts->vecs_sensip[0],ts->numcost,&th->VecsDeltaMu);
1050:   }
1051:   if (ts->vecs_sensi2) {
1052:     VecDuplicateVecs(ts->vecs_sensi[0],ts->numcost,&th->VecsDeltaLam2);
1053:     VecDuplicateVecs(ts->vecs_sensi2[0],ts->numcost,&th->VecsSensi2Temp);
1054:     /* hack ts to make implicit TS solver work when users provide only explicit versions of callbacks (RHSFunction,RHSJacobian,RHSHessian etc.) */
1055:     if (!ts->ihessianproduct_fuu) ts->vecs_fuu = ts->vecs_guu;
1056:     if (!ts->ihessianproduct_fup) ts->vecs_fup = ts->vecs_gup;
1057:   }
1058:   if (ts->vecs_sensi2p) {
1059:     VecDuplicateVecs(ts->vecs_sensi2p[0],ts->numcost,&th->VecsDeltaMu2);
1060:     /* hack ts to make implicit TS solver work when users provide only explicit versions of callbacks (RHSFunction,RHSJacobian,RHSHessian etc.) */
1061:     if (!ts->ihessianproduct_fpu) ts->vecs_fpu = ts->vecs_gpu;
1062:     if (!ts->ihessianproduct_fpp) ts->vecs_fpp = ts->vecs_gpp;
1063:   }
1064:   return(0);
1065: }

1067: static PetscErrorCode TSSetFromOptions_Theta(PetscOptionItems *PetscOptionsObject,TS ts)
1068: {
1069:   TS_Theta       *th = (TS_Theta*)ts->data;

1073:   PetscOptionsHead(PetscOptionsObject,"Theta ODE solver options");
1074:   {
1075:     PetscOptionsReal("-ts_theta_theta","Location of stage (0<Theta<=1)","TSThetaSetTheta",th->Theta,&th->Theta,NULL);
1076:     PetscOptionsBool("-ts_theta_endpoint","Use the endpoint instead of midpoint form of the Theta method","TSThetaSetEndpoint",th->endpoint,&th->endpoint,NULL);
1077:     PetscOptionsBool("-ts_theta_initial_guess_extrapolate","Extrapolate stage initial guess from previous solution (sometimes unstable)","TSThetaSetExtrapolate",th->extrapolate,&th->extrapolate,NULL);
1078:   }
1079:   PetscOptionsTail();
1080:   return(0);
1081: }

1083: static PetscErrorCode TSView_Theta(TS ts,PetscViewer viewer)
1084: {
1085:   TS_Theta       *th = (TS_Theta*)ts->data;
1086:   PetscBool      iascii;

1090:   PetscObjectTypeCompare((PetscObject)viewer,PETSCVIEWERASCII,&iascii);
1091:   if (iascii) {
1092:     PetscViewerASCIIPrintf(viewer,"  Theta=%g\n",(double)th->Theta);
1093:     PetscViewerASCIIPrintf(viewer,"  Extrapolation=%s\n",th->extrapolate ? "yes" : "no");
1094:   }
1095:   return(0);
1096: }

1098: static PetscErrorCode TSThetaGetTheta_Theta(TS ts,PetscReal *theta)
1099: {
1100:   TS_Theta *th = (TS_Theta*)ts->data;

1103:   *theta = th->Theta;
1104:   return(0);
1105: }

1107: static PetscErrorCode TSThetaSetTheta_Theta(TS ts,PetscReal theta)
1108: {
1109:   TS_Theta *th = (TS_Theta*)ts->data;

1112:   if (theta <= 0 || 1 < theta) SETERRQ1(PetscObjectComm((PetscObject)ts),PETSC_ERR_ARG_OUTOFRANGE,"Theta %g not in range (0,1]",(double)theta);
1113:   th->Theta = theta;
1114:   th->order = (th->Theta == 0.5) ? 2 : 1;
1115:   return(0);
1116: }

1118: static PetscErrorCode TSThetaGetEndpoint_Theta(TS ts,PetscBool *endpoint)
1119: {
1120:   TS_Theta *th = (TS_Theta*)ts->data;

1123:   *endpoint = th->endpoint;
1124:   return(0);
1125: }

1127: static PetscErrorCode TSThetaSetEndpoint_Theta(TS ts,PetscBool flg)
1128: {
1129:   TS_Theta *th = (TS_Theta*)ts->data;

1132:   th->endpoint = flg;
1133:   return(0);
1134: }

1136: #if defined(PETSC_HAVE_COMPLEX)
1137: static PetscErrorCode TSComputeLinearStability_Theta(TS ts,PetscReal xr,PetscReal xi,PetscReal *yr,PetscReal *yi)
1138: {
1139:   PetscComplex   z   = xr + xi*PETSC_i,f;
1140:   TS_Theta       *th = (TS_Theta*)ts->data;
1141:   const PetscReal one = 1.0;

1144:   f   = (one + (one - th->Theta)*z)/(one - th->Theta*z);
1145:   *yr = PetscRealPartComplex(f);
1146:   *yi = PetscImaginaryPartComplex(f);
1147:   return(0);
1148: }
1149: #endif

1151: static PetscErrorCode TSGetStages_Theta(TS ts,PetscInt *ns,Vec **Y)
1152: {
1153:   TS_Theta     *th = (TS_Theta*)ts->data;

1156:   if (ns) *ns = 1;
1157:   if (Y)  *Y  = th->endpoint ? &(th->X0) : &(th->X);
1158:   return(0);
1159: }

1161: /* ------------------------------------------------------------ */
1162: /*MC
1163:       TSTHETA - DAE solver using the implicit Theta method

1165:    Level: beginner

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

1172:    Notes:
1173: $  -ts_type theta -ts_theta_theta 1.0 corresponds to backward Euler (TSBEULER)
1174: $  -ts_type theta -ts_theta_theta 0.5 corresponds to the implicit midpoint rule
1175: $  -ts_type theta -ts_theta_theta 0.5 -ts_theta_endpoint corresponds to Crank-Nicholson (TSCN)

1177:    This method can be applied to DAE.

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

1181: .vb
1182:   Theta | Theta
1183:   -------------
1184:         |  1
1185: .ve

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

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

1191: .vb
1192:   0 | 0         0
1193:   1 | 1-Theta   Theta
1194:   -------------------
1195:     | 1-Theta   Theta
1196: .ve

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

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

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

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

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

1208: M*/
1209: PETSC_EXTERN PetscErrorCode TSCreate_Theta(TS ts)
1210: {
1211:   TS_Theta       *th;

1215:   ts->ops->reset           = TSReset_Theta;
1216:   ts->ops->adjointreset    = TSAdjointReset_Theta;
1217:   ts->ops->destroy         = TSDestroy_Theta;
1218:   ts->ops->view            = TSView_Theta;
1219:   ts->ops->setup           = TSSetUp_Theta;
1220:   ts->ops->adjointsetup    = TSAdjointSetUp_Theta;
1221:   ts->ops->adjointreset    = TSAdjointReset_Theta;
1222:   ts->ops->step            = TSStep_Theta;
1223:   ts->ops->interpolate     = TSInterpolate_Theta;
1224:   ts->ops->evaluatewlte    = TSEvaluateWLTE_Theta;
1225:   ts->ops->rollback        = TSRollBack_Theta;
1226:   ts->ops->setfromoptions  = TSSetFromOptions_Theta;
1227:   ts->ops->snesfunction    = SNESTSFormFunction_Theta;
1228:   ts->ops->snesjacobian    = SNESTSFormJacobian_Theta;
1229: #if defined(PETSC_HAVE_COMPLEX)
1230:   ts->ops->linearstability = TSComputeLinearStability_Theta;
1231: #endif
1232:   ts->ops->getstages       = TSGetStages_Theta;
1233:   ts->ops->adjointstep     = TSAdjointStep_Theta;
1234:   ts->ops->adjointintegral = TSAdjointCostIntegral_Theta;
1235:   ts->ops->forwardintegral = TSForwardCostIntegral_Theta;
1236:   ts->default_adapt_type   = TSADAPTNONE;

1238:   ts->ops->forwardsetup     = TSForwardSetUp_Theta;
1239:   ts->ops->forwardreset     = TSForwardReset_Theta;
1240:   ts->ops->forwardstep      = TSForwardStep_Theta;
1241:   ts->ops->forwardgetstages = TSForwardGetStages_Theta;

1243:   ts->usessnes = PETSC_TRUE;

1245:   PetscNewLog(ts,&th);
1246:   ts->data = (void*)th;

1248:   th->VecsDeltaLam    = NULL;
1249:   th->VecsDeltaMu     = NULL;
1250:   th->VecsSensiTemp   = NULL;
1251:   th->VecsSensi2Temp  = NULL;

1253:   th->extrapolate = PETSC_FALSE;
1254:   th->Theta       = 0.5;
1255:   th->order       = 2;
1256:   PetscObjectComposeFunction((PetscObject)ts,"TSThetaGetTheta_C",TSThetaGetTheta_Theta);
1257:   PetscObjectComposeFunction((PetscObject)ts,"TSThetaSetTheta_C",TSThetaSetTheta_Theta);
1258:   PetscObjectComposeFunction((PetscObject)ts,"TSThetaGetEndpoint_C",TSThetaGetEndpoint_Theta);
1259:   PetscObjectComposeFunction((PetscObject)ts,"TSThetaSetEndpoint_C",TSThetaSetEndpoint_Theta);
1260:   return(0);
1261: }

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

1266:   Not Collective

1268:   Input Parameter:
1269: .  ts - timestepping context

1271:   Output Parameter:
1272: .  theta - stage abscissa

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

1277:   Level: Advanced

1279: .seealso: TSThetaSetTheta()
1280: @*/
1281: PetscErrorCode  TSThetaGetTheta(TS ts,PetscReal *theta)
1282: {

1288:   PetscUseMethod(ts,"TSThetaGetTheta_C",(TS,PetscReal*),(ts,theta));
1289:   return(0);
1290: }

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

1295:   Not Collective

1297:   Input Parameter:
1298: +  ts - timestepping context
1299: -  theta - stage abscissa

1301:   Options Database:
1302: .  -ts_theta_theta <theta>

1304:   Level: Intermediate

1306: .seealso: TSThetaGetTheta()
1307: @*/
1308: PetscErrorCode  TSThetaSetTheta(TS ts,PetscReal theta)
1309: {

1314:   PetscTryMethod(ts,"TSThetaSetTheta_C",(TS,PetscReal),(ts,theta));
1315:   return(0);
1316: }

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

1321:   Not Collective

1323:   Input Parameter:
1324: .  ts - timestepping context

1326:   Output Parameter:
1327: .  endpoint - PETSC_TRUE when using the endpoint variant

1329:   Level: Advanced

1331: .seealso: TSThetaSetEndpoint(), TSTHETA, TSCN
1332: @*/
1333: PetscErrorCode TSThetaGetEndpoint(TS ts,PetscBool *endpoint)
1334: {

1340:   PetscUseMethod(ts,"TSThetaGetEndpoint_C",(TS,PetscBool*),(ts,endpoint));
1341:   return(0);
1342: }

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

1347:   Not Collective

1349:   Input Parameter:
1350: +  ts - timestepping context
1351: -  flg - PETSC_TRUE to use the endpoint variant

1353:   Options Database:
1354: .  -ts_theta_endpoint <flg>

1356:   Level: Intermediate

1358: .seealso: TSTHETA, TSCN
1359: @*/
1360: PetscErrorCode TSThetaSetEndpoint(TS ts,PetscBool flg)
1361: {

1366:   PetscTryMethod(ts,"TSThetaSetEndpoint_C",(TS,PetscBool),(ts,flg));
1367:   return(0);
1368: }

1370: /*
1371:  * TSBEULER and TSCN are straightforward specializations of TSTHETA.
1372:  * The creation functions for these specializations are below.
1373:  */

1375: static PetscErrorCode TSSetUp_BEuler(TS ts)
1376: {
1377:   TS_Theta       *th = (TS_Theta*)ts->data;

1381:   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");
1382:   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");
1383:   TSSetUp_Theta(ts);
1384:   return(0);
1385: }

1387: static PetscErrorCode TSView_BEuler(TS ts,PetscViewer viewer)
1388: {
1390:   return(0);
1391: }

1393: /*MC
1394:       TSBEULER - ODE solver using the implicit backward Euler method

1396:   Level: beginner

1398:   Notes:
1399:   TSBEULER is equivalent to TSTHETA with Theta=1.0

1401: $  -ts_type theta -ts_theta_theta 1.0

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

1405: M*/
1406: PETSC_EXTERN PetscErrorCode TSCreate_BEuler(TS ts)
1407: {

1411:   TSCreate_Theta(ts);
1412:   TSThetaSetTheta(ts,1.0);
1413:   TSThetaSetEndpoint(ts,PETSC_FALSE);
1414:   ts->ops->setup = TSSetUp_BEuler;
1415:   ts->ops->view  = TSView_BEuler;
1416:   return(0);
1417: }

1419: static PetscErrorCode TSSetUp_CN(TS ts)
1420: {
1421:   TS_Theta       *th = (TS_Theta*)ts->data;

1425:   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");
1426:   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");
1427:   TSSetUp_Theta(ts);
1428:   return(0);
1429: }

1431: static PetscErrorCode TSView_CN(TS ts,PetscViewer viewer)
1432: {
1434:   return(0);
1435: }

1437: /*MC
1438:       TSCN - ODE solver using the implicit Crank-Nicolson method.

1440:   Level: beginner

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

1445: $  -ts_type theta -ts_theta_theta 0.5 -ts_theta_endpoint

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

1449: M*/
1450: PETSC_EXTERN PetscErrorCode TSCreate_CN(TS ts)
1451: {

1455:   TSCreate_Theta(ts);
1456:   TSThetaSetTheta(ts,0.5);
1457:   TSThetaSetEndpoint(ts,PETSC_TRUE);
1458:   ts->ops->setup = TSSetUp_CN;
1459:   ts->ops->view  = TSView_CN;
1460:   return(0);
1461: }