Actual source code: theta.c

  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          Stages[2];   /* Storage for stage solutions */
 13:   Vec          X0, X, Xdot; /* Storage for u^n, u^n + dt a_{11} k_1, and time derivative u^{n+1}_t */
 14:   Vec          affine;      /* Affine vector needed for residual at beginning of step in endpoint formulation */
 15:   PetscReal    Theta;
 16:   PetscReal    shift; /* Shift parameter for SNES Jacobian, used by forward, TLM and adjoint */
 17:   PetscInt     order;
 18:   PetscBool    endpoint;
 19:   PetscBool    extrapolate;
 20:   TSStepStatus status;
 21:   Vec          VecCostIntegral0; /* Backup for roll-backs due to events, used by cost integral */
 22:   PetscReal    ptime0;           /* Backup for ts->ptime, the start time of current time step, used by TLM and cost integral */
 23:   PetscReal    time_step0;       /* Backup for ts->timestep, the step size of current time step, used by TLM and cost integral*/

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

 45: static PetscErrorCode TSThetaGetX0AndXdot(TS ts, DM dm, Vec *X0, Vec *Xdot)
 46: {
 47:   TS_Theta *th = (TS_Theta *)ts->data;

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

 63: static PetscErrorCode TSThetaRestoreX0AndXdot(TS ts, DM dm, Vec *X0, Vec *Xdot)
 64: {
 65:   PetscFunctionBegin;
 66:   if (X0) {
 67:     if (dm && dm != ts->dm) PetscCall(DMRestoreNamedGlobalVector(dm, "TSTheta_X0", X0));
 68:   }
 69:   if (Xdot) {
 70:     if (dm && dm != ts->dm) PetscCall(DMRestoreNamedGlobalVector(dm, "TSTheta_Xdot", Xdot));
 71:   }
 72:   PetscFunctionReturn(PETSC_SUCCESS);
 73: }

 75: static PetscErrorCode DMCoarsenHook_TSTheta(DM fine, DM coarse, void *ctx)
 76: {
 77:   PetscFunctionBegin;
 78:   PetscFunctionReturn(PETSC_SUCCESS);
 79: }

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

 86:   PetscFunctionBegin;
 87:   PetscCall(TSThetaGetX0AndXdot(ts, fine, &X0, &Xdot));
 88:   PetscCall(TSThetaGetX0AndXdot(ts, coarse, &X0_c, &Xdot_c));
 89:   PetscCall(MatRestrict(restrct, X0, X0_c));
 90:   PetscCall(MatRestrict(restrct, Xdot, Xdot_c));
 91:   PetscCall(VecPointwiseMult(X0_c, rscale, X0_c));
 92:   PetscCall(VecPointwiseMult(Xdot_c, rscale, Xdot_c));
 93:   PetscCall(TSThetaRestoreX0AndXdot(ts, fine, &X0, &Xdot));
 94:   PetscCall(TSThetaRestoreX0AndXdot(ts, coarse, &X0_c, &Xdot_c));
 95:   PetscFunctionReturn(PETSC_SUCCESS);
 96: }

 98: static PetscErrorCode DMSubDomainHook_TSTheta(DM dm, DM subdm, void *ctx)
 99: {
100:   PetscFunctionBegin;
101:   PetscFunctionReturn(PETSC_SUCCESS);
102: }

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

109:   PetscFunctionBegin;
110:   PetscCall(TSThetaGetX0AndXdot(ts, dm, &X0, &Xdot));
111:   PetscCall(TSThetaGetX0AndXdot(ts, subdm, &X0_sub, &Xdot_sub));

113:   PetscCall(VecScatterBegin(gscat, X0, X0_sub, INSERT_VALUES, SCATTER_FORWARD));
114:   PetscCall(VecScatterEnd(gscat, X0, X0_sub, INSERT_VALUES, SCATTER_FORWARD));

116:   PetscCall(VecScatterBegin(gscat, Xdot, Xdot_sub, INSERT_VALUES, SCATTER_FORWARD));
117:   PetscCall(VecScatterEnd(gscat, Xdot, Xdot_sub, INSERT_VALUES, SCATTER_FORWARD));

119:   PetscCall(TSThetaRestoreX0AndXdot(ts, dm, &X0, &Xdot));
120:   PetscCall(TSThetaRestoreX0AndXdot(ts, subdm, &X0_sub, &Xdot_sub));
121:   PetscFunctionReturn(PETSC_SUCCESS);
122: }

124: static PetscErrorCode TSThetaEvaluateCostIntegral(TS ts)
125: {
126:   TS_Theta *th     = (TS_Theta *)ts->data;
127:   TS        quadts = ts->quadraturets;

129:   PetscFunctionBegin;
130:   if (th->endpoint) {
131:     /* Evolve ts->vec_costintegral to compute integrals */
132:     if (th->Theta != 1.0) {
133:       PetscCall(TSComputeRHSFunction(quadts, th->ptime0, th->X0, ts->vec_costintegrand));
134:       PetscCall(VecAXPY(quadts->vec_sol, th->time_step0 * (1.0 - th->Theta), ts->vec_costintegrand));
135:     }
136:     PetscCall(TSComputeRHSFunction(quadts, ts->ptime, ts->vec_sol, ts->vec_costintegrand));
137:     PetscCall(VecAXPY(quadts->vec_sol, th->time_step0 * th->Theta, ts->vec_costintegrand));
138:   } else {
139:     PetscCall(TSComputeRHSFunction(quadts, th->stage_time, th->X, ts->vec_costintegrand));
140:     PetscCall(VecAXPY(quadts->vec_sol, th->time_step0, ts->vec_costintegrand));
141:   }
142:   PetscFunctionReturn(PETSC_SUCCESS);
143: }

145: static PetscErrorCode TSForwardCostIntegral_Theta(TS ts)
146: {
147:   TS_Theta *th     = (TS_Theta *)ts->data;
148:   TS        quadts = ts->quadraturets;

150:   PetscFunctionBegin;
151:   /* backup cost integral */
152:   PetscCall(VecCopy(quadts->vec_sol, th->VecCostIntegral0));
153:   PetscCall(TSThetaEvaluateCostIntegral(ts));
154:   PetscFunctionReturn(PETSC_SUCCESS);
155: }

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

161:   PetscFunctionBegin;
162:   /* Like TSForwardCostIntegral(), the adjoint cost integral evaluation relies on ptime0 and time_step0. */
163:   th->ptime0     = ts->ptime + ts->time_step;
164:   th->time_step0 = -ts->time_step;
165:   PetscCall(TSThetaEvaluateCostIntegral(ts));
166:   PetscFunctionReturn(PETSC_SUCCESS);
167: }

169: static PetscErrorCode TSTheta_SNESSolve(TS ts, Vec b, Vec x)
170: {
171:   PetscInt nits, lits;

173:   PetscFunctionBegin;
174:   PetscCall(SNESSolve(ts->snes, b, x));
175:   PetscCall(SNESGetIterationNumber(ts->snes, &nits));
176:   PetscCall(SNESGetLinearSolveIterations(ts->snes, &lits));
177:   ts->snes_its += nits;
178:   ts->ksp_its += lits;
179:   PetscFunctionReturn(PETSC_SUCCESS);
180: }

182: static PetscErrorCode TSResizeRegister_Theta(TS ts, PetscBool reg)
183: {
184:   TS_Theta *th = (TS_Theta *)ts->data;

186:   PetscFunctionBegin;
187:   if (reg) {
188:     PetscCall(TSResizeRegisterVec(ts, "ts:theta:sol_prev", th->vec_sol_prev));
189:     PetscCall(TSResizeRegisterVec(ts, "ts:theta:X0", th->X0));
190:   } else {
191:     PetscCall(TSResizeRetrieveVec(ts, "ts:theta:sol_prev", &th->vec_sol_prev));
192:     PetscCall(PetscObjectReference((PetscObject)th->vec_sol_prev));
193:     PetscCall(TSResizeRetrieveVec(ts, "ts:theta:X0", &th->X0));
194:     PetscCall(PetscObjectReference((PetscObject)th->X0));
195:   }
196:   PetscFunctionReturn(PETSC_SUCCESS);
197: }

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

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

212:   th->status = TS_STEP_INCOMPLETE;
213:   while (!ts->reason && th->status != TS_STEP_COMPLETE) {
214:     th->shift      = 1 / (th->Theta * ts->time_step);
215:     th->stage_time = ts->ptime + (th->endpoint ? (PetscReal)1 : th->Theta) * ts->time_step;
216:     PetscCall(VecCopy(th->X0, th->X));
217:     if (th->extrapolate && !ts->steprestart) PetscCall(VecAXPY(th->X, 1 / th->shift, th->Xdot));
218:     if (th->endpoint) { /* This formulation assumes linear time-independent mass matrix */
219:       if (!th->affine) PetscCall(VecDuplicate(ts->vec_sol, &th->affine));
220:       PetscCall(VecZeroEntries(th->Xdot));
221:       PetscCall(TSComputeIFunction(ts, ts->ptime, th->X0, th->Xdot, th->affine, PETSC_FALSE));
222:       PetscCall(VecScale(th->affine, (th->Theta - 1) / th->Theta));
223:     }
224:     PetscCall(TSPreStage(ts, th->stage_time));
225:     PetscCall(TSTheta_SNESSolve(ts, th->endpoint ? th->affine : NULL, th->X));
226:     PetscCall(TSPostStage(ts, th->stage_time, 0, &th->X));
227:     PetscCall(TSAdaptCheckStage(ts->adapt, ts, th->stage_time, th->X, &stageok));
228:     if (!stageok) goto reject_step;

230:     if (th->endpoint) {
231:       PetscCall(VecCopy(th->X, ts->vec_sol));
232:     } else {
233:       PetscCall(VecAXPBYPCZ(th->Xdot, -th->shift, th->shift, 0, th->X0, th->X)); /* th->Xdot is needed by TSInterpolate_Theta */
234:       if (th->Theta == 1.0) PetscCall(VecCopy(th->X, ts->vec_sol));              /* BEULER, stage already checked */
235:       else {
236:         PetscCall(VecAXPY(ts->vec_sol, ts->time_step, th->Xdot));
237:         PetscCall(TSAdaptCheckStage(ts->adapt, ts, ts->ptime + ts->time_step, ts->vec_sol, &stageok));
238:         if (!stageok) {
239:           PetscCall(VecCopy(th->X0, ts->vec_sol));
240:           goto reject_step;
241:         }
242:       }
243:     }

245:     th->status = TS_STEP_PENDING;
246:     PetscCall(TSAdaptChoose(ts->adapt, ts, ts->time_step, NULL, &next_time_step, &accept));
247:     th->status = accept ? TS_STEP_COMPLETE : TS_STEP_INCOMPLETE;
248:     if (!accept) {
249:       PetscCall(VecCopy(th->X0, ts->vec_sol));
250:       ts->time_step = next_time_step;
251:       goto reject_step;
252:     }

254:     if (ts->forward_solve || ts->costintegralfwd) { /* Save the info for the later use in cost integral evaluation */
255:       th->ptime0     = ts->ptime;
256:       th->time_step0 = ts->time_step;
257:     }
258:     ts->ptime += ts->time_step;
259:     ts->time_step = next_time_step;
260:     break;

262:   reject_step:
263:     ts->reject++;
264:     accept = PETSC_FALSE;
265:     if (!ts->reason && ++rejections > ts->max_reject && ts->max_reject >= 0) {
266:       ts->reason = TS_DIVERGED_STEP_REJECTED;
267:       PetscCall(PetscInfo(ts, "Step=%" PetscInt_FMT ", step rejections %" PetscInt_FMT " greater than current TS allowed, stopping solve\n", ts->steps, rejections));
268:     }
269:   }
270:   PetscFunctionReturn(PETSC_SUCCESS);
271: }

273: static PetscErrorCode TSAdjointStepBEuler_Private(TS ts)
274: {
275:   TS_Theta      *th           = (TS_Theta *)ts->data;
276:   TS             quadts       = ts->quadraturets;
277:   Vec           *VecsDeltaLam = th->VecsDeltaLam, *VecsDeltaMu = th->VecsDeltaMu, *VecsSensiTemp = th->VecsSensiTemp;
278:   Vec           *VecsDeltaLam2 = th->VecsDeltaLam2, *VecsDeltaMu2 = th->VecsDeltaMu2, *VecsSensi2Temp = th->VecsSensi2Temp;
279:   PetscInt       nadj;
280:   Mat            J, Jpre, quadJ = NULL, quadJp = NULL;
281:   KSP            ksp;
282:   PetscScalar   *xarr;
283:   TSEquationType eqtype;
284:   PetscBool      isexplicitode = PETSC_FALSE;
285:   PetscReal      adjoint_time_step;

287:   PetscFunctionBegin;
288:   PetscCall(TSGetEquationType(ts, &eqtype));
289:   if (eqtype == TS_EQ_ODE_EXPLICIT) {
290:     isexplicitode = PETSC_TRUE;
291:     VecsDeltaLam  = ts->vecs_sensi;
292:     VecsDeltaLam2 = ts->vecs_sensi2;
293:   }
294:   th->status = TS_STEP_INCOMPLETE;
295:   PetscCall(SNESGetKSP(ts->snes, &ksp));
296:   PetscCall(TSGetIJacobian(ts, &J, &Jpre, NULL, NULL));
297:   if (quadts) {
298:     PetscCall(TSGetRHSJacobian(quadts, &quadJ, NULL, NULL, NULL));
299:     PetscCall(TSGetRHSJacobianP(quadts, &quadJp, NULL, NULL));
300:   }

302:   th->stage_time    = ts->ptime;
303:   adjoint_time_step = -ts->time_step; /* always positive since time_step is negative */

305:   /* Build RHS for first-order adjoint lambda_{n+1}/h + r_u^T(n+1) */
306:   if (quadts) PetscCall(TSComputeRHSJacobian(quadts, th->stage_time, ts->vec_sol, quadJ, NULL));

308:   for (nadj = 0; nadj < ts->numcost; nadj++) {
309:     PetscCall(VecCopy(ts->vecs_sensi[nadj], VecsSensiTemp[nadj]));
310:     PetscCall(VecScale(VecsSensiTemp[nadj], 1. / adjoint_time_step)); /* lambda_{n+1}/h */
311:     if (quadJ) {
312:       PetscCall(MatDenseGetColumn(quadJ, nadj, &xarr));
313:       PetscCall(VecPlaceArray(ts->vec_drdu_col, xarr));
314:       PetscCall(VecAXPY(VecsSensiTemp[nadj], 1., ts->vec_drdu_col));
315:       PetscCall(VecResetArray(ts->vec_drdu_col));
316:       PetscCall(MatDenseRestoreColumn(quadJ, &xarr));
317:     }
318:   }

320:   /* Build LHS for first-order adjoint */
321:   th->shift = 1. / adjoint_time_step;
322:   PetscCall(TSComputeSNESJacobian(ts, ts->vec_sol, J, Jpre));
323:   PetscCall(KSPSetOperators(ksp, J, Jpre));

325:   /* Solve stage equation LHS*lambda_s = RHS for first-order adjoint */
326:   for (nadj = 0; nadj < ts->numcost; nadj++) {
327:     KSPConvergedReason kspreason;
328:     PetscCall(KSPSolveTranspose(ksp, VecsSensiTemp[nadj], VecsDeltaLam[nadj]));
329:     PetscCall(KSPGetConvergedReason(ksp, &kspreason));
330:     if (kspreason < 0) {
331:       ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
332:       PetscCall(PetscInfo(ts, "Step=%" PetscInt_FMT ", %" PetscInt_FMT "th cost function, transposed linear solve fails, stopping 1st-order adjoint solve\n", ts->steps, nadj));
333:     }
334:   }

336:   if (ts->vecs_sensi2) { /* U_{n+1} */
337:     /* Get w1 at t_{n+1} from TLM matrix */
338:     PetscCall(MatDenseGetColumn(ts->mat_sensip, 0, &xarr));
339:     PetscCall(VecPlaceArray(ts->vec_sensip_col, xarr));
340:     /* lambda_s^T F_UU w_1 */
341:     PetscCall(TSComputeIHessianProductFunctionUU(ts, th->stage_time, ts->vec_sol, VecsDeltaLam, ts->vec_sensip_col, ts->vecs_fuu));
342:     /* lambda_s^T F_UP w_2 */
343:     PetscCall(TSComputeIHessianProductFunctionUP(ts, th->stage_time, ts->vec_sol, VecsDeltaLam, ts->vec_dir, ts->vecs_fup));
344:     for (nadj = 0; nadj < ts->numcost; nadj++) { /* compute the residual */
345:       PetscCall(VecCopy(ts->vecs_sensi2[nadj], VecsSensi2Temp[nadj]));
346:       PetscCall(VecScale(VecsSensi2Temp[nadj], 1. / adjoint_time_step));
347:       PetscCall(VecAXPY(VecsSensi2Temp[nadj], -1., ts->vecs_fuu[nadj]));
348:       if (ts->vecs_fup) PetscCall(VecAXPY(VecsSensi2Temp[nadj], -1., ts->vecs_fup[nadj]));
349:     }
350:     /* Solve stage equation LHS X = RHS for second-order adjoint */
351:     for (nadj = 0; nadj < ts->numcost; nadj++) {
352:       KSPConvergedReason kspreason;
353:       PetscCall(KSPSolveTranspose(ksp, VecsSensi2Temp[nadj], VecsDeltaLam2[nadj]));
354:       PetscCall(KSPGetConvergedReason(ksp, &kspreason));
355:       if (kspreason < 0) {
356:         ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
357:         PetscCall(PetscInfo(ts, "Step=%" PetscInt_FMT ", %" PetscInt_FMT "th cost function, transposed linear solve fails, stopping 2nd-order adjoint solve\n", ts->steps, nadj));
358:       }
359:     }
360:   }

362:   /* Update sensitivities, and evaluate integrals if there is any */
363:   if (!isexplicitode) {
364:     th->shift = 0.0;
365:     PetscCall(TSComputeSNESJacobian(ts, ts->vec_sol, J, Jpre));
366:     PetscCall(KSPSetOperators(ksp, J, Jpre));
367:     for (nadj = 0; nadj < ts->numcost; nadj++) {
368:       /* Add f_U \lambda_s to the original RHS */
369:       PetscCall(VecScale(VecsSensiTemp[nadj], -1.));
370:       PetscCall(MatMultTransposeAdd(J, VecsDeltaLam[nadj], VecsSensiTemp[nadj], VecsSensiTemp[nadj]));
371:       PetscCall(VecScale(VecsSensiTemp[nadj], -adjoint_time_step));
372:       PetscCall(VecCopy(VecsSensiTemp[nadj], ts->vecs_sensi[nadj]));
373:       if (ts->vecs_sensi2) {
374:         PetscCall(MatMultTransposeAdd(J, VecsDeltaLam2[nadj], VecsSensi2Temp[nadj], VecsSensi2Temp[nadj]));
375:         PetscCall(VecScale(VecsSensi2Temp[nadj], -adjoint_time_step));
376:         PetscCall(VecCopy(VecsSensi2Temp[nadj], ts->vecs_sensi2[nadj]));
377:       }
378:     }
379:   }
380:   if (ts->vecs_sensip) {
381:     PetscCall(VecAXPBYPCZ(th->Xdot, -1. / adjoint_time_step, 1.0 / adjoint_time_step, 0, th->X0, ts->vec_sol));
382:     PetscCall(TSComputeIJacobianP(ts, th->stage_time, ts->vec_sol, th->Xdot, 1. / adjoint_time_step, ts->Jacp, PETSC_FALSE)); /* get -f_p */
383:     if (quadts) PetscCall(TSComputeRHSJacobianP(quadts, th->stage_time, ts->vec_sol, quadJp));
384:     if (ts->vecs_sensi2p) {
385:       /* lambda_s^T F_PU w_1 */
386:       PetscCall(TSComputeIHessianProductFunctionPU(ts, th->stage_time, ts->vec_sol, VecsDeltaLam, ts->vec_sensip_col, ts->vecs_fpu));
387:       /* lambda_s^T F_PP w_2 */
388:       PetscCall(TSComputeIHessianProductFunctionPP(ts, th->stage_time, ts->vec_sol, VecsDeltaLam, ts->vec_dir, ts->vecs_fpp));
389:     }

391:     for (nadj = 0; nadj < ts->numcost; nadj++) {
392:       PetscCall(MatMultTranspose(ts->Jacp, VecsDeltaLam[nadj], VecsDeltaMu[nadj]));
393:       PetscCall(VecAXPY(ts->vecs_sensip[nadj], -adjoint_time_step, VecsDeltaMu[nadj]));
394:       if (quadJp) {
395:         PetscCall(MatDenseGetColumn(quadJp, nadj, &xarr));
396:         PetscCall(VecPlaceArray(ts->vec_drdp_col, xarr));
397:         PetscCall(VecAXPY(ts->vecs_sensip[nadj], adjoint_time_step, ts->vec_drdp_col));
398:         PetscCall(VecResetArray(ts->vec_drdp_col));
399:         PetscCall(MatDenseRestoreColumn(quadJp, &xarr));
400:       }
401:       if (ts->vecs_sensi2p) {
402:         PetscCall(MatMultTranspose(ts->Jacp, VecsDeltaLam2[nadj], VecsDeltaMu2[nadj]));
403:         PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step, VecsDeltaMu2[nadj]));
404:         if (ts->vecs_fpu) PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step, ts->vecs_fpu[nadj]));
405:         if (ts->vecs_fpp) PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step, ts->vecs_fpp[nadj]));
406:       }
407:     }
408:   }

410:   if (ts->vecs_sensi2) {
411:     PetscCall(VecResetArray(ts->vec_sensip_col));
412:     PetscCall(MatDenseRestoreColumn(ts->mat_sensip, &xarr));
413:   }
414:   th->status = TS_STEP_COMPLETE;
415:   PetscFunctionReturn(PETSC_SUCCESS);
416: }

418: static PetscErrorCode TSAdjointStep_Theta(TS ts)
419: {
420:   TS_Theta    *th           = (TS_Theta *)ts->data;
421:   TS           quadts       = ts->quadraturets;
422:   Vec         *VecsDeltaLam = th->VecsDeltaLam, *VecsDeltaMu = th->VecsDeltaMu, *VecsSensiTemp = th->VecsSensiTemp;
423:   Vec         *VecsDeltaLam2 = th->VecsDeltaLam2, *VecsDeltaMu2 = th->VecsDeltaMu2, *VecsSensi2Temp = th->VecsSensi2Temp;
424:   PetscInt     nadj;
425:   Mat          J, Jpre, quadJ = NULL, quadJp = NULL;
426:   KSP          ksp;
427:   PetscScalar *xarr;
428:   PetscReal    adjoint_time_step;
429:   PetscReal    adjoint_ptime; /* end time of the adjoint time step (ts->ptime is the start time, usually ts->ptime is larger than adjoint_ptime) */

431:   PetscFunctionBegin;
432:   if (th->Theta == 1.) {
433:     PetscCall(TSAdjointStepBEuler_Private(ts));
434:     PetscFunctionReturn(PETSC_SUCCESS);
435:   }
436:   th->status = TS_STEP_INCOMPLETE;
437:   PetscCall(SNESGetKSP(ts->snes, &ksp));
438:   PetscCall(TSGetIJacobian(ts, &J, &Jpre, NULL, NULL));
439:   if (quadts) {
440:     PetscCall(TSGetRHSJacobian(quadts, &quadJ, NULL, NULL, NULL));
441:     PetscCall(TSGetRHSJacobianP(quadts, &quadJp, NULL, NULL));
442:   }
443:   /* If endpoint=1, th->ptime and th->X0 will be used; if endpoint=0, th->stage_time and th->X will be used. */
444:   th->stage_time    = th->endpoint ? ts->ptime : (ts->ptime + (1. - th->Theta) * ts->time_step);
445:   adjoint_ptime     = ts->ptime + ts->time_step;
446:   adjoint_time_step = -ts->time_step; /* always positive since time_step is negative */

448:   if (!th->endpoint) {
449:     /* recover th->X0 using vec_sol and the stage value th->X */
450:     PetscCall(VecAXPBYPCZ(th->X0, 1.0 / (1.0 - th->Theta), th->Theta / (th->Theta - 1.0), 0, th->X, ts->vec_sol));
451:   }

453:   /* Build RHS for first-order adjoint */
454:   /* Cost function has an integral term */
455:   if (quadts) {
456:     if (th->endpoint) {
457:       PetscCall(TSComputeRHSJacobian(quadts, th->stage_time, ts->vec_sol, quadJ, NULL));
458:     } else {
459:       PetscCall(TSComputeRHSJacobian(quadts, th->stage_time, th->X, quadJ, NULL));
460:     }
461:   }

463:   for (nadj = 0; nadj < ts->numcost; nadj++) {
464:     PetscCall(VecCopy(ts->vecs_sensi[nadj], VecsSensiTemp[nadj]));
465:     PetscCall(VecScale(VecsSensiTemp[nadj], 1. / (th->Theta * adjoint_time_step)));
466:     if (quadJ) {
467:       PetscCall(MatDenseGetColumn(quadJ, nadj, &xarr));
468:       PetscCall(VecPlaceArray(ts->vec_drdu_col, xarr));
469:       PetscCall(VecAXPY(VecsSensiTemp[nadj], 1., ts->vec_drdu_col));
470:       PetscCall(VecResetArray(ts->vec_drdu_col));
471:       PetscCall(MatDenseRestoreColumn(quadJ, &xarr));
472:     }
473:   }

475:   /* Build LHS for first-order adjoint */
476:   th->shift = 1. / (th->Theta * adjoint_time_step);
477:   if (th->endpoint) {
478:     PetscCall(TSComputeSNESJacobian(ts, ts->vec_sol, J, Jpre));
479:   } else {
480:     PetscCall(TSComputeSNESJacobian(ts, th->X, J, Jpre));
481:   }
482:   PetscCall(KSPSetOperators(ksp, J, Jpre));

484:   /* Solve stage equation LHS*lambda_s = RHS for first-order adjoint */
485:   for (nadj = 0; nadj < ts->numcost; nadj++) {
486:     KSPConvergedReason kspreason;
487:     PetscCall(KSPSolveTranspose(ksp, VecsSensiTemp[nadj], VecsDeltaLam[nadj]));
488:     PetscCall(KSPGetConvergedReason(ksp, &kspreason));
489:     if (kspreason < 0) {
490:       ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
491:       PetscCall(PetscInfo(ts, "Step=%" PetscInt_FMT ", %" PetscInt_FMT "th cost function, transposed linear solve fails, stopping 1st-order adjoint solve\n", ts->steps, nadj));
492:     }
493:   }

495:   /* Second-order adjoint */
496:   if (ts->vecs_sensi2) { /* U_{n+1} */
497:     PetscCheck(th->endpoint, PetscObjectComm((PetscObject)ts), PETSC_ERR_SUP, "Operation not implemented in TS_Theta");
498:     /* Get w1 at t_{n+1} from TLM matrix */
499:     PetscCall(MatDenseGetColumn(ts->mat_sensip, 0, &xarr));
500:     PetscCall(VecPlaceArray(ts->vec_sensip_col, xarr));
501:     /* lambda_s^T F_UU w_1 */
502:     PetscCall(TSComputeIHessianProductFunctionUU(ts, th->stage_time, ts->vec_sol, VecsDeltaLam, ts->vec_sensip_col, ts->vecs_fuu));
503:     PetscCall(VecResetArray(ts->vec_sensip_col));
504:     PetscCall(MatDenseRestoreColumn(ts->mat_sensip, &xarr));
505:     /* lambda_s^T F_UP w_2 */
506:     PetscCall(TSComputeIHessianProductFunctionUP(ts, th->stage_time, ts->vec_sol, VecsDeltaLam, ts->vec_dir, ts->vecs_fup));
507:     for (nadj = 0; nadj < ts->numcost; nadj++) { /* compute the residual */
508:       PetscCall(VecCopy(ts->vecs_sensi2[nadj], VecsSensi2Temp[nadj]));
509:       PetscCall(VecScale(VecsSensi2Temp[nadj], th->shift));
510:       PetscCall(VecAXPY(VecsSensi2Temp[nadj], -1., ts->vecs_fuu[nadj]));
511:       if (ts->vecs_fup) PetscCall(VecAXPY(VecsSensi2Temp[nadj], -1., ts->vecs_fup[nadj]));
512:     }
513:     /* Solve stage equation LHS X = RHS for second-order adjoint */
514:     for (nadj = 0; nadj < ts->numcost; nadj++) {
515:       KSPConvergedReason kspreason;
516:       PetscCall(KSPSolveTranspose(ksp, VecsSensi2Temp[nadj], VecsDeltaLam2[nadj]));
517:       PetscCall(KSPGetConvergedReason(ksp, &kspreason));
518:       if (kspreason < 0) {
519:         ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
520:         PetscCall(PetscInfo(ts, "Step=%" PetscInt_FMT ", %" PetscInt_FMT "th cost function, transposed linear solve fails, stopping 2nd-order adjoint solve\n", ts->steps, nadj));
521:       }
522:     }
523:   }

525:   /* Update sensitivities, and evaluate integrals if there is any */
526:   if (th->endpoint) { /* two-stage Theta methods with th->Theta!=1, th->Theta==1 leads to BEuler */
527:     th->shift      = 1. / ((th->Theta - 1.) * adjoint_time_step);
528:     th->stage_time = adjoint_ptime;
529:     PetscCall(TSComputeSNESJacobian(ts, th->X0, J, Jpre));
530:     PetscCall(KSPSetOperators(ksp, J, Jpre));
531:     /* R_U at t_n */
532:     if (quadts) PetscCall(TSComputeRHSJacobian(quadts, adjoint_ptime, th->X0, quadJ, NULL));
533:     for (nadj = 0; nadj < ts->numcost; nadj++) {
534:       PetscCall(MatMultTranspose(J, VecsDeltaLam[nadj], ts->vecs_sensi[nadj]));
535:       if (quadJ) {
536:         PetscCall(MatDenseGetColumn(quadJ, nadj, &xarr));
537:         PetscCall(VecPlaceArray(ts->vec_drdu_col, xarr));
538:         PetscCall(VecAXPY(ts->vecs_sensi[nadj], -1., ts->vec_drdu_col));
539:         PetscCall(VecResetArray(ts->vec_drdu_col));
540:         PetscCall(MatDenseRestoreColumn(quadJ, &xarr));
541:       }
542:       PetscCall(VecScale(ts->vecs_sensi[nadj], 1. / th->shift));
543:     }

545:     /* Second-order adjoint */
546:     if (ts->vecs_sensi2) { /* U_n */
547:       /* Get w1 at t_n from TLM matrix */
548:       PetscCall(MatDenseGetColumn(th->MatFwdSensip0, 0, &xarr));
549:       PetscCall(VecPlaceArray(ts->vec_sensip_col, xarr));
550:       /* lambda_s^T F_UU w_1 */
551:       PetscCall(TSComputeIHessianProductFunctionUU(ts, adjoint_ptime, th->X0, VecsDeltaLam, ts->vec_sensip_col, ts->vecs_fuu));
552:       PetscCall(VecResetArray(ts->vec_sensip_col));
553:       PetscCall(MatDenseRestoreColumn(th->MatFwdSensip0, &xarr));
554:       /* lambda_s^T F_UU w_2 */
555:       PetscCall(TSComputeIHessianProductFunctionUP(ts, adjoint_ptime, th->X0, VecsDeltaLam, ts->vec_dir, ts->vecs_fup));
556:       for (nadj = 0; nadj < ts->numcost; nadj++) {
557:         /* 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  */
558:         PetscCall(MatMultTranspose(J, VecsDeltaLam2[nadj], ts->vecs_sensi2[nadj]));
559:         PetscCall(VecAXPY(ts->vecs_sensi2[nadj], 1., ts->vecs_fuu[nadj]));
560:         if (ts->vecs_fup) PetscCall(VecAXPY(ts->vecs_sensi2[nadj], 1., ts->vecs_fup[nadj]));
561:         PetscCall(VecScale(ts->vecs_sensi2[nadj], 1. / th->shift));
562:       }
563:     }

565:     th->stage_time = ts->ptime; /* recover the old value */

567:     if (ts->vecs_sensip) { /* sensitivities wrt parameters */
568:       /* U_{n+1} */
569:       th->shift = 1.0 / (adjoint_time_step * th->Theta);
570:       PetscCall(VecAXPBYPCZ(th->Xdot, -th->shift, th->shift, 0, th->X0, ts->vec_sol));
571:       PetscCall(TSComputeIJacobianP(ts, th->stage_time, ts->vec_sol, th->Xdot, -1. / (th->Theta * adjoint_time_step), ts->Jacp, PETSC_FALSE));
572:       if (quadts) PetscCall(TSComputeRHSJacobianP(quadts, th->stage_time, ts->vec_sol, quadJp));
573:       for (nadj = 0; nadj < ts->numcost; nadj++) {
574:         PetscCall(MatMultTranspose(ts->Jacp, VecsDeltaLam[nadj], VecsDeltaMu[nadj]));
575:         PetscCall(VecAXPY(ts->vecs_sensip[nadj], -adjoint_time_step * th->Theta, VecsDeltaMu[nadj]));
576:         if (quadJp) {
577:           PetscCall(MatDenseGetColumn(quadJp, nadj, &xarr));
578:           PetscCall(VecPlaceArray(ts->vec_drdp_col, xarr));
579:           PetscCall(VecAXPY(ts->vecs_sensip[nadj], adjoint_time_step * th->Theta, ts->vec_drdp_col));
580:           PetscCall(VecResetArray(ts->vec_drdp_col));
581:           PetscCall(MatDenseRestoreColumn(quadJp, &xarr));
582:         }
583:       }
584:       if (ts->vecs_sensi2p) { /* second-order */
585:         /* Get w1 at t_{n+1} from TLM matrix */
586:         PetscCall(MatDenseGetColumn(ts->mat_sensip, 0, &xarr));
587:         PetscCall(VecPlaceArray(ts->vec_sensip_col, xarr));
588:         /* lambda_s^T F_PU w_1 */
589:         PetscCall(TSComputeIHessianProductFunctionPU(ts, th->stage_time, ts->vec_sol, VecsDeltaLam, ts->vec_sensip_col, ts->vecs_fpu));
590:         PetscCall(VecResetArray(ts->vec_sensip_col));
591:         PetscCall(MatDenseRestoreColumn(ts->mat_sensip, &xarr));

593:         /* lambda_s^T F_PP w_2 */
594:         PetscCall(TSComputeIHessianProductFunctionPP(ts, th->stage_time, ts->vec_sol, VecsDeltaLam, ts->vec_dir, ts->vecs_fpp));
595:         for (nadj = 0; nadj < ts->numcost; nadj++) {
596:           /* 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)  */
597:           PetscCall(MatMultTranspose(ts->Jacp, VecsDeltaLam2[nadj], VecsDeltaMu2[nadj]));
598:           PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step * th->Theta, VecsDeltaMu2[nadj]));
599:           if (ts->vecs_fpu) PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step * th->Theta, ts->vecs_fpu[nadj]));
600:           if (ts->vecs_fpp) PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step * th->Theta, ts->vecs_fpp[nadj]));
601:         }
602:       }

604:       /* U_s */
605:       PetscCall(VecZeroEntries(th->Xdot));
606:       PetscCall(TSComputeIJacobianP(ts, adjoint_ptime, th->X0, th->Xdot, 1. / ((th->Theta - 1.0) * adjoint_time_step), ts->Jacp, PETSC_FALSE));
607:       if (quadts) PetscCall(TSComputeRHSJacobianP(quadts, adjoint_ptime, th->X0, quadJp));
608:       for (nadj = 0; nadj < ts->numcost; nadj++) {
609:         PetscCall(MatMultTranspose(ts->Jacp, VecsDeltaLam[nadj], VecsDeltaMu[nadj]));
610:         PetscCall(VecAXPY(ts->vecs_sensip[nadj], -adjoint_time_step * (1.0 - th->Theta), VecsDeltaMu[nadj]));
611:         if (quadJp) {
612:           PetscCall(MatDenseGetColumn(quadJp, nadj, &xarr));
613:           PetscCall(VecPlaceArray(ts->vec_drdp_col, xarr));
614:           PetscCall(VecAXPY(ts->vecs_sensip[nadj], adjoint_time_step * (1.0 - th->Theta), ts->vec_drdp_col));
615:           PetscCall(VecResetArray(ts->vec_drdp_col));
616:           PetscCall(MatDenseRestoreColumn(quadJp, &xarr));
617:         }
618:         if (ts->vecs_sensi2p) { /* second-order */
619:           /* Get w1 at t_n from TLM matrix */
620:           PetscCall(MatDenseGetColumn(th->MatFwdSensip0, 0, &xarr));
621:           PetscCall(VecPlaceArray(ts->vec_sensip_col, xarr));
622:           /* lambda_s^T F_PU w_1 */
623:           PetscCall(TSComputeIHessianProductFunctionPU(ts, adjoint_ptime, th->X0, VecsDeltaLam, ts->vec_sensip_col, ts->vecs_fpu));
624:           PetscCall(VecResetArray(ts->vec_sensip_col));
625:           PetscCall(MatDenseRestoreColumn(th->MatFwdSensip0, &xarr));
626:           /* lambda_s^T F_PP w_2 */
627:           PetscCall(TSComputeIHessianProductFunctionPP(ts, adjoint_ptime, th->X0, VecsDeltaLam, ts->vec_dir, ts->vecs_fpp));
628:           for (nadj = 0; nadj < ts->numcost; nadj++) {
629:             /* 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) */
630:             PetscCall(MatMultTranspose(ts->Jacp, VecsDeltaLam2[nadj], VecsDeltaMu2[nadj]));
631:             PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step * (1.0 - th->Theta), VecsDeltaMu2[nadj]));
632:             if (ts->vecs_fpu) PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step * (1.0 - th->Theta), ts->vecs_fpu[nadj]));
633:             if (ts->vecs_fpp) PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step * (1.0 - th->Theta), ts->vecs_fpp[nadj]));
634:           }
635:         }
636:       }
637:     }
638:   } else { /* one-stage case */
639:     th->shift = 0.0;
640:     PetscCall(TSComputeSNESJacobian(ts, th->X, J, Jpre)); /* get -f_y */
641:     PetscCall(KSPSetOperators(ksp, J, Jpre));
642:     if (quadts) PetscCall(TSComputeRHSJacobian(quadts, th->stage_time, th->X, quadJ, NULL));
643:     for (nadj = 0; nadj < ts->numcost; nadj++) {
644:       PetscCall(MatMultTranspose(J, VecsDeltaLam[nadj], VecsSensiTemp[nadj]));
645:       PetscCall(VecAXPY(ts->vecs_sensi[nadj], -adjoint_time_step, VecsSensiTemp[nadj]));
646:       if (quadJ) {
647:         PetscCall(MatDenseGetColumn(quadJ, nadj, &xarr));
648:         PetscCall(VecPlaceArray(ts->vec_drdu_col, xarr));
649:         PetscCall(VecAXPY(ts->vecs_sensi[nadj], adjoint_time_step, ts->vec_drdu_col));
650:         PetscCall(VecResetArray(ts->vec_drdu_col));
651:         PetscCall(MatDenseRestoreColumn(quadJ, &xarr));
652:       }
653:     }
654:     if (ts->vecs_sensip) {
655:       th->shift = 1.0 / (adjoint_time_step * th->Theta);
656:       PetscCall(VecAXPBYPCZ(th->Xdot, -th->shift, th->shift, 0, th->X0, th->X));
657:       PetscCall(TSComputeIJacobianP(ts, th->stage_time, th->X, th->Xdot, th->shift, ts->Jacp, PETSC_FALSE));
658:       if (quadts) PetscCall(TSComputeRHSJacobianP(quadts, th->stage_time, th->X, quadJp));
659:       for (nadj = 0; nadj < ts->numcost; nadj++) {
660:         PetscCall(MatMultTranspose(ts->Jacp, VecsDeltaLam[nadj], VecsDeltaMu[nadj]));
661:         PetscCall(VecAXPY(ts->vecs_sensip[nadj], -adjoint_time_step, VecsDeltaMu[nadj]));
662:         if (quadJp) {
663:           PetscCall(MatDenseGetColumn(quadJp, nadj, &xarr));
664:           PetscCall(VecPlaceArray(ts->vec_drdp_col, xarr));
665:           PetscCall(VecAXPY(ts->vecs_sensip[nadj], adjoint_time_step, ts->vec_drdp_col));
666:           PetscCall(VecResetArray(ts->vec_drdp_col));
667:           PetscCall(MatDenseRestoreColumn(quadJp, &xarr));
668:         }
669:       }
670:     }
671:   }

673:   th->status = TS_STEP_COMPLETE;
674:   PetscFunctionReturn(PETSC_SUCCESS);
675: }

677: static PetscErrorCode TSInterpolate_Theta(TS ts, PetscReal t, Vec X)
678: {
679:   TS_Theta *th = (TS_Theta *)ts->data;
680:   PetscReal dt = t - ts->ptime;

682:   PetscFunctionBegin;
683:   PetscCall(VecCopy(ts->vec_sol, th->X));
684:   if (th->endpoint) dt *= th->Theta;
685:   PetscCall(VecWAXPY(X, dt, th->Xdot, th->X));
686:   PetscFunctionReturn(PETSC_SUCCESS);
687: }

689: static PetscErrorCode TSEvaluateWLTE_Theta(TS ts, NormType wnormtype, PetscInt *order, PetscReal *wlte)
690: {
691:   TS_Theta *th = (TS_Theta *)ts->data;
692:   Vec       X  = ts->vec_sol;      /* X = solution */
693:   Vec       Y  = th->vec_lte_work; /* Y = X + LTE  */
694:   PetscReal wltea, wlter;

696:   PetscFunctionBegin;
697:   if (!th->vec_sol_prev) {
698:     *wlte = -1;
699:     PetscFunctionReturn(PETSC_SUCCESS);
700:   }
701:   /* Cannot compute LTE in first step or in restart after event */
702:   if (ts->steprestart) {
703:     *wlte = -1;
704:     PetscFunctionReturn(PETSC_SUCCESS);
705:   }
706:   /* Compute LTE using backward differences with non-constant time step */
707:   {
708:     PetscReal   h = ts->time_step, h_prev = ts->ptime - ts->ptime_prev;
709:     PetscReal   a = 1 + h_prev / h;
710:     PetscScalar scal[3];
711:     Vec         vecs[3];
712:     scal[0] = -1 / a;
713:     scal[1] = +1 / (a - 1);
714:     scal[2] = -1 / (a * (a - 1));
715:     vecs[0] = X;
716:     vecs[1] = th->X0;
717:     vecs[2] = th->vec_sol_prev;
718:     PetscCall(VecCopy(X, Y));
719:     PetscCall(VecMAXPY(Y, 3, scal, vecs));
720:     PetscCall(TSErrorWeightedNorm(ts, X, Y, wnormtype, wlte, &wltea, &wlter));
721:   }
722:   if (order) *order = 2;
723:   PetscFunctionReturn(PETSC_SUCCESS);
724: }

726: static PetscErrorCode TSRollBack_Theta(TS ts)
727: {
728:   TS_Theta *th     = (TS_Theta *)ts->data;
729:   TS        quadts = ts->quadraturets;

731:   PetscFunctionBegin;
732:   if (quadts && ts->costintegralfwd) PetscCall(VecCopy(th->VecCostIntegral0, quadts->vec_sol));
733:   th->status = TS_STEP_INCOMPLETE;
734:   if (ts->mat_sensip) PetscCall(MatCopy(th->MatFwdSensip0, ts->mat_sensip, SAME_NONZERO_PATTERN));
735:   if (quadts && quadts->mat_sensip) PetscCall(MatCopy(th->MatIntegralSensip0, quadts->mat_sensip, SAME_NONZERO_PATTERN));
736:   PetscFunctionReturn(PETSC_SUCCESS);
737: }

739: static PetscErrorCode TSForwardStep_Theta(TS ts)
740: {
741:   TS_Theta    *th                   = (TS_Theta *)ts->data;
742:   TS           quadts               = ts->quadraturets;
743:   Mat          MatDeltaFwdSensip    = th->MatDeltaFwdSensip;
744:   Vec          VecDeltaFwdSensipCol = th->VecDeltaFwdSensipCol;
745:   PetscInt     ntlm;
746:   KSP          ksp;
747:   Mat          J, Jpre, quadJ = NULL, quadJp = NULL;
748:   PetscScalar *barr, *xarr;
749:   PetscReal    previous_shift;

751:   PetscFunctionBegin;
752:   previous_shift = th->shift;
753:   PetscCall(MatCopy(ts->mat_sensip, th->MatFwdSensip0, SAME_NONZERO_PATTERN));

755:   if (quadts && quadts->mat_sensip) PetscCall(MatCopy(quadts->mat_sensip, th->MatIntegralSensip0, SAME_NONZERO_PATTERN));
756:   PetscCall(SNESGetKSP(ts->snes, &ksp));
757:   PetscCall(TSGetIJacobian(ts, &J, &Jpre, NULL, NULL));
758:   if (quadts) {
759:     PetscCall(TSGetRHSJacobian(quadts, &quadJ, NULL, NULL, NULL));
760:     PetscCall(TSGetRHSJacobianP(quadts, &quadJp, NULL, NULL));
761:   }

763:   /* Build RHS */
764:   if (th->endpoint) { /* 2-stage method*/
765:     th->shift = 1. / ((th->Theta - 1.) * th->time_step0);
766:     PetscCall(TSComputeIJacobian(ts, th->ptime0, th->X0, th->Xdot, th->shift, J, Jpre, PETSC_FALSE));
767:     PetscCall(MatMatMult(J, ts->mat_sensip, MAT_REUSE_MATRIX, PETSC_DETERMINE, &MatDeltaFwdSensip));
768:     PetscCall(MatScale(MatDeltaFwdSensip, (th->Theta - 1.) / th->Theta));

770:     /* Add the f_p forcing terms */
771:     if (ts->Jacp) {
772:       PetscCall(VecZeroEntries(th->Xdot));
773:       PetscCall(TSComputeIJacobianP(ts, th->ptime0, th->X0, th->Xdot, th->shift, ts->Jacp, PETSC_FALSE));
774:       PetscCall(MatAXPY(MatDeltaFwdSensip, (th->Theta - 1.) / th->Theta, ts->Jacp, SUBSET_NONZERO_PATTERN));
775:       th->shift = previous_shift;
776:       PetscCall(VecAXPBYPCZ(th->Xdot, -th->shift, th->shift, 0, th->X0, ts->vec_sol));
777:       PetscCall(TSComputeIJacobianP(ts, th->stage_time, ts->vec_sol, th->Xdot, th->shift, ts->Jacp, PETSC_FALSE));
778:       PetscCall(MatAXPY(MatDeltaFwdSensip, -1., ts->Jacp, SUBSET_NONZERO_PATTERN));
779:     }
780:   } else { /* 1-stage method */
781:     th->shift = 0.0;
782:     PetscCall(TSComputeIJacobian(ts, th->stage_time, th->X, th->Xdot, th->shift, J, Jpre, PETSC_FALSE));
783:     PetscCall(MatMatMult(J, ts->mat_sensip, MAT_REUSE_MATRIX, PETSC_DETERMINE, &MatDeltaFwdSensip));
784:     PetscCall(MatScale(MatDeltaFwdSensip, -1.));

786:     /* Add the f_p forcing terms */
787:     if (ts->Jacp) {
788:       th->shift = previous_shift;
789:       PetscCall(VecAXPBYPCZ(th->Xdot, -th->shift, th->shift, 0, th->X0, th->X));
790:       PetscCall(TSComputeIJacobianP(ts, th->stage_time, th->X, th->Xdot, th->shift, ts->Jacp, PETSC_FALSE));
791:       PetscCall(MatAXPY(MatDeltaFwdSensip, -1., ts->Jacp, SUBSET_NONZERO_PATTERN));
792:     }
793:   }

795:   /* Build LHS */
796:   th->shift = previous_shift; /* recover the previous shift used in TSStep_Theta() */
797:   if (th->endpoint) {
798:     PetscCall(TSComputeIJacobian(ts, th->stage_time, ts->vec_sol, th->Xdot, th->shift, J, Jpre, PETSC_FALSE));
799:   } else {
800:     PetscCall(TSComputeIJacobian(ts, th->stage_time, th->X, th->Xdot, th->shift, J, Jpre, PETSC_FALSE));
801:   }
802:   PetscCall(KSPSetOperators(ksp, J, Jpre));

804:   /*
805:     Evaluate the first stage of integral gradients with the 2-stage method:
806:     drdu|t_n*S(t_n) + drdp|t_n
807:     This is done before the linear solve because the sensitivity variable S(t_n) will be propagated to S(t_{n+1})
808:   */
809:   if (th->endpoint) { /* 2-stage method only */
810:     if (quadts && quadts->mat_sensip) {
811:       PetscCall(TSComputeRHSJacobian(quadts, th->ptime0, th->X0, quadJ, NULL));
812:       PetscCall(TSComputeRHSJacobianP(quadts, th->ptime0, th->X0, quadJp));
813:       PetscCall(MatTransposeMatMult(ts->mat_sensip, quadJ, MAT_REUSE_MATRIX, PETSC_DETERMINE, &th->MatIntegralSensipTemp));
814:       PetscCall(MatAXPY(th->MatIntegralSensipTemp, 1, quadJp, SAME_NONZERO_PATTERN));
815:       PetscCall(MatAXPY(quadts->mat_sensip, th->time_step0 * (1. - th->Theta), th->MatIntegralSensipTemp, SAME_NONZERO_PATTERN));
816:     }
817:   }

819:   /* Solve the tangent linear equation for forward sensitivities to parameters */
820:   for (ntlm = 0; ntlm < th->num_tlm; ntlm++) {
821:     KSPConvergedReason kspreason;
822:     PetscCall(MatDenseGetColumn(MatDeltaFwdSensip, ntlm, &barr));
823:     PetscCall(VecPlaceArray(VecDeltaFwdSensipCol, barr));
824:     if (th->endpoint) {
825:       PetscCall(MatDenseGetColumn(ts->mat_sensip, ntlm, &xarr));
826:       PetscCall(VecPlaceArray(ts->vec_sensip_col, xarr));
827:       PetscCall(KSPSolve(ksp, VecDeltaFwdSensipCol, ts->vec_sensip_col));
828:       PetscCall(VecResetArray(ts->vec_sensip_col));
829:       PetscCall(MatDenseRestoreColumn(ts->mat_sensip, &xarr));
830:     } else {
831:       PetscCall(KSPSolve(ksp, VecDeltaFwdSensipCol, VecDeltaFwdSensipCol));
832:     }
833:     PetscCall(KSPGetConvergedReason(ksp, &kspreason));
834:     if (kspreason < 0) {
835:       ts->reason = TSFORWARD_DIVERGED_LINEAR_SOLVE;
836:       PetscCall(PetscInfo(ts, "Step=%" PetscInt_FMT ", %" PetscInt_FMT "th tangent linear solve, linear solve fails, stopping tangent linear solve\n", ts->steps, ntlm));
837:     }
838:     PetscCall(VecResetArray(VecDeltaFwdSensipCol));
839:     PetscCall(MatDenseRestoreColumn(MatDeltaFwdSensip, &barr));
840:   }

842:   /*
843:     Evaluate the second stage of integral gradients with the 2-stage method:
844:     drdu|t_{n+1}*S(t_{n+1}) + drdp|t_{n+1}
845:   */
846:   if (quadts && quadts->mat_sensip) {
847:     if (!th->endpoint) {
848:       PetscCall(MatAXPY(ts->mat_sensip, 1, MatDeltaFwdSensip, SAME_NONZERO_PATTERN)); /* stage sensitivity */
849:       PetscCall(TSComputeRHSJacobian(quadts, th->stage_time, th->X, quadJ, NULL));
850:       PetscCall(TSComputeRHSJacobianP(quadts, th->stage_time, th->X, quadJp));
851:       PetscCall(MatTransposeMatMult(ts->mat_sensip, quadJ, MAT_REUSE_MATRIX, PETSC_DETERMINE, &th->MatIntegralSensipTemp));
852:       PetscCall(MatAXPY(th->MatIntegralSensipTemp, 1, quadJp, SAME_NONZERO_PATTERN));
853:       PetscCall(MatAXPY(quadts->mat_sensip, th->time_step0, th->MatIntegralSensipTemp, SAME_NONZERO_PATTERN));
854:       PetscCall(MatAXPY(ts->mat_sensip, (1. - th->Theta) / th->Theta, MatDeltaFwdSensip, SAME_NONZERO_PATTERN));
855:     } else {
856:       PetscCall(TSComputeRHSJacobian(quadts, th->stage_time, ts->vec_sol, quadJ, NULL));
857:       PetscCall(TSComputeRHSJacobianP(quadts, th->stage_time, ts->vec_sol, quadJp));
858:       PetscCall(MatTransposeMatMult(ts->mat_sensip, quadJ, MAT_REUSE_MATRIX, PETSC_DETERMINE, &th->MatIntegralSensipTemp));
859:       PetscCall(MatAXPY(th->MatIntegralSensipTemp, 1, quadJp, SAME_NONZERO_PATTERN));
860:       PetscCall(MatAXPY(quadts->mat_sensip, th->time_step0 * th->Theta, th->MatIntegralSensipTemp, SAME_NONZERO_PATTERN));
861:     }
862:   } else {
863:     if (!th->endpoint) PetscCall(MatAXPY(ts->mat_sensip, 1. / th->Theta, MatDeltaFwdSensip, SAME_NONZERO_PATTERN));
864:   }
865:   PetscFunctionReturn(PETSC_SUCCESS);
866: }

868: static PetscErrorCode TSForwardGetStages_Theta(TS ts, PetscInt *ns, Mat *stagesensip[])
869: {
870:   TS_Theta *th = (TS_Theta *)ts->data;

872:   PetscFunctionBegin;
873:   if (ns) {
874:     if (!th->endpoint && th->Theta != 1.0) *ns = 1; /* midpoint form */
875:     else *ns = 2;                                   /* endpoint form */
876:   }
877:   if (stagesensip) {
878:     if (!th->endpoint && th->Theta != 1.0) {
879:       th->MatFwdStages[0] = th->MatDeltaFwdSensip;
880:     } else {
881:       th->MatFwdStages[0] = th->MatFwdSensip0;
882:       th->MatFwdStages[1] = ts->mat_sensip; /* stiffly accurate */
883:     }
884:     *stagesensip = th->MatFwdStages;
885:   }
886:   PetscFunctionReturn(PETSC_SUCCESS);
887: }

889: /*------------------------------------------------------------*/
890: static PetscErrorCode TSReset_Theta(TS ts)
891: {
892:   TS_Theta *th = (TS_Theta *)ts->data;

894:   PetscFunctionBegin;
895:   PetscCall(VecDestroy(&th->X));
896:   PetscCall(VecDestroy(&th->Xdot));
897:   PetscCall(VecDestroy(&th->X0));
898:   PetscCall(VecDestroy(&th->affine));

900:   PetscCall(VecDestroy(&th->vec_sol_prev));
901:   PetscCall(VecDestroy(&th->vec_lte_work));

903:   PetscCall(VecDestroy(&th->VecCostIntegral0));
904:   PetscFunctionReturn(PETSC_SUCCESS);
905: }

907: static PetscErrorCode TSAdjointReset_Theta(TS ts)
908: {
909:   TS_Theta *th = (TS_Theta *)ts->data;

911:   PetscFunctionBegin;
912:   PetscCall(VecDestroyVecs(ts->numcost, &th->VecsDeltaLam));
913:   PetscCall(VecDestroyVecs(ts->numcost, &th->VecsDeltaMu));
914:   PetscCall(VecDestroyVecs(ts->numcost, &th->VecsDeltaLam2));
915:   PetscCall(VecDestroyVecs(ts->numcost, &th->VecsDeltaMu2));
916:   PetscCall(VecDestroyVecs(ts->numcost, &th->VecsSensiTemp));
917:   PetscCall(VecDestroyVecs(ts->numcost, &th->VecsSensi2Temp));
918:   PetscFunctionReturn(PETSC_SUCCESS);
919: }

921: static PetscErrorCode TSDestroy_Theta(TS ts)
922: {
923:   PetscFunctionBegin;
924:   PetscCall(TSReset_Theta(ts));
925:   if (ts->dm) {
926:     PetscCall(DMCoarsenHookRemove(ts->dm, DMCoarsenHook_TSTheta, DMRestrictHook_TSTheta, ts));
927:     PetscCall(DMSubDomainHookRemove(ts->dm, DMSubDomainHook_TSTheta, DMSubDomainRestrictHook_TSTheta, ts));
928:   }
929:   PetscCall(PetscFree(ts->data));
930:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaGetTheta_C", NULL));
931:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaSetTheta_C", NULL));
932:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaGetEndpoint_C", NULL));
933:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaSetEndpoint_C", NULL));
934:   PetscFunctionReturn(PETSC_SUCCESS);
935: }

937: /*
938:   This defines the nonlinear equation that is to be solved with SNES
939:   G(U) = F[t0+Theta*dt, U, (U-U0)*shift] = 0

941:   Note that U here is the stage argument. This means that U = U_{n+1} only if endpoint = true,
942:   otherwise U = theta U_{n+1} + (1 - theta) U0, which for the case of implicit midpoint is
943:   U = (U_{n+1} + U0)/2
944: */
945: static PetscErrorCode SNESTSFormFunction_Theta(SNES snes, Vec x, Vec y, TS ts)
946: {
947:   TS_Theta *th = (TS_Theta *)ts->data;
948:   Vec       X0, Xdot;
949:   DM        dm, dmsave;
950:   PetscReal shift = th->shift;

952:   PetscFunctionBegin;
953:   PetscCall(SNESGetDM(snes, &dm));
954:   /* When using the endpoint variant, this is actually 1/Theta * Xdot */
955:   PetscCall(TSThetaGetX0AndXdot(ts, dm, &X0, &Xdot));
956:   if (x != X0) {
957:     PetscCall(VecAXPBYPCZ(Xdot, -shift, shift, 0, X0, x));
958:   } else {
959:     PetscCall(VecZeroEntries(Xdot));
960:   }
961:   /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */
962:   dmsave = ts->dm;
963:   ts->dm = dm;
964:   PetscCall(TSComputeIFunction(ts, th->stage_time, x, Xdot, y, PETSC_FALSE));
965:   ts->dm = dmsave;
966:   PetscCall(TSThetaRestoreX0AndXdot(ts, dm, &X0, &Xdot));
967:   PetscFunctionReturn(PETSC_SUCCESS);
968: }

970: static PetscErrorCode SNESTSFormJacobian_Theta(SNES snes, Vec x, Mat A, Mat B, TS ts)
971: {
972:   TS_Theta *th = (TS_Theta *)ts->data;
973:   Vec       Xdot;
974:   DM        dm, dmsave;
975:   PetscReal shift = th->shift;

977:   PetscFunctionBegin;
978:   PetscCall(SNESGetDM(snes, &dm));
979:   /* Xdot has already been computed in SNESTSFormFunction_Theta (SNES guarantees this) */
980:   PetscCall(TSThetaGetX0AndXdot(ts, dm, NULL, &Xdot));

982:   dmsave = ts->dm;
983:   ts->dm = dm;
984:   PetscCall(TSComputeIJacobian(ts, th->stage_time, x, Xdot, shift, A, B, PETSC_FALSE));
985:   ts->dm = dmsave;
986:   PetscCall(TSThetaRestoreX0AndXdot(ts, dm, NULL, &Xdot));
987:   PetscFunctionReturn(PETSC_SUCCESS);
988: }

990: static PetscErrorCode TSForwardSetUp_Theta(TS ts)
991: {
992:   TS_Theta *th     = (TS_Theta *)ts->data;
993:   TS        quadts = ts->quadraturets;

995:   PetscFunctionBegin;
996:   /* combine sensitivities to parameters and sensitivities to initial values into one array */
997:   th->num_tlm = ts->num_parameters;
998:   PetscCall(MatDuplicate(ts->mat_sensip, MAT_DO_NOT_COPY_VALUES, &th->MatDeltaFwdSensip));
999:   if (quadts && quadts->mat_sensip) {
1000:     PetscCall(MatDuplicate(quadts->mat_sensip, MAT_DO_NOT_COPY_VALUES, &th->MatIntegralSensipTemp));
1001:     PetscCall(MatDuplicate(quadts->mat_sensip, MAT_DO_NOT_COPY_VALUES, &th->MatIntegralSensip0));
1002:   }
1003:   /* backup sensitivity results for roll-backs */
1004:   PetscCall(MatDuplicate(ts->mat_sensip, MAT_DO_NOT_COPY_VALUES, &th->MatFwdSensip0));

1006:   PetscCall(VecDuplicate(ts->vec_sol, &th->VecDeltaFwdSensipCol));
1007:   PetscFunctionReturn(PETSC_SUCCESS);
1008: }

1010: static PetscErrorCode TSForwardReset_Theta(TS ts)
1011: {
1012:   TS_Theta *th     = (TS_Theta *)ts->data;
1013:   TS        quadts = ts->quadraturets;

1015:   PetscFunctionBegin;
1016:   if (quadts && quadts->mat_sensip) {
1017:     PetscCall(MatDestroy(&th->MatIntegralSensipTemp));
1018:     PetscCall(MatDestroy(&th->MatIntegralSensip0));
1019:   }
1020:   PetscCall(VecDestroy(&th->VecDeltaFwdSensipCol));
1021:   PetscCall(MatDestroy(&th->MatDeltaFwdSensip));
1022:   PetscCall(MatDestroy(&th->MatFwdSensip0));
1023:   PetscFunctionReturn(PETSC_SUCCESS);
1024: }

1026: static PetscErrorCode TSSetUp_Theta(TS ts)
1027: {
1028:   TS_Theta *th     = (TS_Theta *)ts->data;
1029:   TS        quadts = ts->quadraturets;
1030:   PetscBool match;

1032:   PetscFunctionBegin;
1033:   if (!th->VecCostIntegral0 && quadts && ts->costintegralfwd) { /* back up cost integral */
1034:     PetscCall(VecDuplicate(quadts->vec_sol, &th->VecCostIntegral0));
1035:   }
1036:   if (!th->X) PetscCall(VecDuplicate(ts->vec_sol, &th->X));
1037:   if (!th->Xdot) PetscCall(VecDuplicate(ts->vec_sol, &th->Xdot));
1038:   if (!th->X0) PetscCall(VecDuplicate(ts->vec_sol, &th->X0));
1039:   if (th->endpoint) PetscCall(VecDuplicate(ts->vec_sol, &th->affine));

1041:   th->order = (th->Theta == 0.5) ? 2 : 1;
1042:   th->shift = 1 / (th->Theta * ts->time_step);

1044:   PetscCall(TSGetDM(ts, &ts->dm));
1045:   PetscCall(DMCoarsenHookAdd(ts->dm, DMCoarsenHook_TSTheta, DMRestrictHook_TSTheta, ts));
1046:   PetscCall(DMSubDomainHookAdd(ts->dm, DMSubDomainHook_TSTheta, DMSubDomainRestrictHook_TSTheta, ts));

1048:   PetscCall(TSGetAdapt(ts, &ts->adapt));
1049:   PetscCall(TSAdaptCandidatesClear(ts->adapt));
1050:   PetscCall(PetscObjectTypeCompare((PetscObject)ts->adapt, TSADAPTNONE, &match));
1051:   if (!match) {
1052:     if (!th->vec_sol_prev) PetscCall(VecDuplicate(ts->vec_sol, &th->vec_sol_prev));
1053:     if (!th->vec_lte_work) PetscCall(VecDuplicate(ts->vec_sol, &th->vec_lte_work));
1054:   }
1055:   PetscCall(TSGetSNES(ts, &ts->snes));

1057:   ts->stifflyaccurate = (!th->endpoint && th->Theta != 1.0) ? PETSC_FALSE : PETSC_TRUE;
1058:   PetscFunctionReturn(PETSC_SUCCESS);
1059: }

1061: /*------------------------------------------------------------*/

1063: static PetscErrorCode TSAdjointSetUp_Theta(TS ts)
1064: {
1065:   TS_Theta *th = (TS_Theta *)ts->data;

1067:   PetscFunctionBegin;
1068:   PetscCall(VecDuplicateVecs(ts->vecs_sensi[0], ts->numcost, &th->VecsDeltaLam));
1069:   PetscCall(VecDuplicateVecs(ts->vecs_sensi[0], ts->numcost, &th->VecsSensiTemp));
1070:   if (ts->vecs_sensip) PetscCall(VecDuplicateVecs(ts->vecs_sensip[0], ts->numcost, &th->VecsDeltaMu));
1071:   if (ts->vecs_sensi2) {
1072:     PetscCall(VecDuplicateVecs(ts->vecs_sensi[0], ts->numcost, &th->VecsDeltaLam2));
1073:     PetscCall(VecDuplicateVecs(ts->vecs_sensi2[0], ts->numcost, &th->VecsSensi2Temp));
1074:     /* hack ts to make implicit TS solver work when users provide only explicit versions of callbacks (RHSFunction,RHSJacobian,RHSHessian etc.) */
1075:     if (!ts->ihessianproduct_fuu) ts->vecs_fuu = ts->vecs_guu;
1076:     if (!ts->ihessianproduct_fup) ts->vecs_fup = ts->vecs_gup;
1077:   }
1078:   if (ts->vecs_sensi2p) {
1079:     PetscCall(VecDuplicateVecs(ts->vecs_sensi2p[0], ts->numcost, &th->VecsDeltaMu2));
1080:     /* hack ts to make implicit TS solver work when users provide only explicit versions of callbacks (RHSFunction,RHSJacobian,RHSHessian etc.) */
1081:     if (!ts->ihessianproduct_fpu) ts->vecs_fpu = ts->vecs_gpu;
1082:     if (!ts->ihessianproduct_fpp) ts->vecs_fpp = ts->vecs_gpp;
1083:   }
1084:   PetscFunctionReturn(PETSC_SUCCESS);
1085: }

1087: static PetscErrorCode TSSetFromOptions_Theta(TS ts, PetscOptionItems PetscOptionsObject)
1088: {
1089:   TS_Theta *th = (TS_Theta *)ts->data;

1091:   PetscFunctionBegin;
1092:   PetscOptionsHeadBegin(PetscOptionsObject, "Theta ODE solver options");
1093:   {
1094:     PetscCall(PetscOptionsReal("-ts_theta_theta", "Location of stage (0<Theta<=1)", "TSThetaSetTheta", th->Theta, &th->Theta, NULL));
1095:     PetscCall(PetscOptionsBool("-ts_theta_endpoint", "Use the endpoint instead of midpoint form of the Theta method", "TSThetaSetEndpoint", th->endpoint, &th->endpoint, NULL));
1096:     PetscCall(PetscOptionsBool("-ts_theta_initial_guess_extrapolate", "Extrapolate stage initial guess from previous solution (sometimes unstable)", "TSThetaSetExtrapolate", th->extrapolate, &th->extrapolate, NULL));
1097:   }
1098:   PetscOptionsHeadEnd();
1099:   PetscFunctionReturn(PETSC_SUCCESS);
1100: }

1102: static PetscErrorCode TSView_Theta(TS ts, PetscViewer viewer)
1103: {
1104:   TS_Theta *th = (TS_Theta *)ts->data;
1105:   PetscBool iascii;

1107:   PetscFunctionBegin;
1108:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &iascii));
1109:   if (iascii) {
1110:     PetscCall(PetscViewerASCIIPrintf(viewer, "  Theta=%g\n", (double)th->Theta));
1111:     PetscCall(PetscViewerASCIIPrintf(viewer, "  Extrapolation=%s\n", th->extrapolate ? "yes" : "no"));
1112:   }
1113:   PetscFunctionReturn(PETSC_SUCCESS);
1114: }

1116: static PetscErrorCode TSThetaGetTheta_Theta(TS ts, PetscReal *theta)
1117: {
1118:   TS_Theta *th = (TS_Theta *)ts->data;

1120:   PetscFunctionBegin;
1121:   *theta = th->Theta;
1122:   PetscFunctionReturn(PETSC_SUCCESS);
1123: }

1125: static PetscErrorCode TSThetaSetTheta_Theta(TS ts, PetscReal theta)
1126: {
1127:   TS_Theta *th = (TS_Theta *)ts->data;

1129:   PetscFunctionBegin;
1130:   PetscCheck(theta > 0 && theta <= 1, PetscObjectComm((PetscObject)ts), PETSC_ERR_ARG_OUTOFRANGE, "Theta %g not in range (0,1]", (double)theta);
1131:   th->Theta = theta;
1132:   th->order = (th->Theta == 0.5) ? 2 : 1;
1133:   PetscFunctionReturn(PETSC_SUCCESS);
1134: }

1136: static PetscErrorCode TSThetaGetEndpoint_Theta(TS ts, PetscBool *endpoint)
1137: {
1138:   TS_Theta *th = (TS_Theta *)ts->data;

1140:   PetscFunctionBegin;
1141:   *endpoint = th->endpoint;
1142:   PetscFunctionReturn(PETSC_SUCCESS);
1143: }

1145: static PetscErrorCode TSThetaSetEndpoint_Theta(TS ts, PetscBool flg)
1146: {
1147:   TS_Theta *th = (TS_Theta *)ts->data;

1149:   PetscFunctionBegin;
1150:   th->endpoint = flg;
1151:   PetscFunctionReturn(PETSC_SUCCESS);
1152: }

1154: #if defined(PETSC_HAVE_COMPLEX)
1155: static PetscErrorCode TSComputeLinearStability_Theta(TS ts, PetscReal xr, PetscReal xi, PetscReal *yr, PetscReal *yi)
1156: {
1157:   PetscComplex z  = xr + xi * PETSC_i, f;
1158:   TS_Theta    *th = (TS_Theta *)ts->data;

1160:   PetscFunctionBegin;
1161:   f   = (1.0 + (1.0 - th->Theta) * z) / (1.0 - th->Theta * z);
1162:   *yr = PetscRealPartComplex(f);
1163:   *yi = PetscImaginaryPartComplex(f);
1164:   PetscFunctionReturn(PETSC_SUCCESS);
1165: }
1166: #endif

1168: static PetscErrorCode TSGetStages_Theta(TS ts, PetscInt *ns, Vec *Y[])
1169: {
1170:   TS_Theta *th = (TS_Theta *)ts->data;

1172:   PetscFunctionBegin;
1173:   if (ns) {
1174:     if (!th->endpoint && th->Theta != 1.0) *ns = 1; /* midpoint form */
1175:     else *ns = 2;                                   /* endpoint form */
1176:   }
1177:   if (Y) {
1178:     if (!th->endpoint && th->Theta != 1.0) {
1179:       th->Stages[0] = th->X;
1180:     } else {
1181:       th->Stages[0] = th->X0;
1182:       th->Stages[1] = ts->vec_sol; /* stiffly accurate */
1183:     }
1184:     *Y = th->Stages;
1185:   }
1186:   PetscFunctionReturn(PETSC_SUCCESS);
1187: }

1189: /* ------------------------------------------------------------ */
1190: /*MC
1191:       TSTHETA - DAE solver using the implicit Theta method

1193:    Level: beginner

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

1200:    Notes:
1201: .vb
1202:   -ts_type theta -ts_theta_theta 1.0 corresponds to backward Euler (TSBEULER)
1203:   -ts_type theta -ts_theta_theta 0.5 corresponds to the implicit midpoint rule
1204:   -ts_type theta -ts_theta_theta 0.5 -ts_theta_endpoint corresponds to Crank-Nicholson (TSCN)
1205: .ve

1207:    The endpoint variant of the Theta method and backward Euler can be applied to DAE. The midpoint variant is not suitable for DAEs because it is not stiffly accurate.

1209:    The midpoint variant is cast as a 1-stage implicit Runge-Kutta method.

1211: .vb
1212:   Theta | Theta
1213:   -------------
1214:         |  1
1215: .ve

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

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

1221: .vb
1222:   0 | 0         0
1223:   1 | 1-Theta   Theta
1224:   -------------------
1225:     | 1-Theta   Theta
1226: .ve

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

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

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

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

1236: .seealso: [](ch_ts), `TSCreate()`, `TS`, `TSSetType()`, `TSCN`, `TSBEULER`, `TSThetaSetTheta()`, `TSThetaSetEndpoint()`
1237: M*/
1238: PETSC_EXTERN PetscErrorCode TSCreate_Theta(TS ts)
1239: {
1240:   TS_Theta *th;

1242:   PetscFunctionBegin;
1243:   ts->ops->reset          = TSReset_Theta;
1244:   ts->ops->adjointreset   = TSAdjointReset_Theta;
1245:   ts->ops->destroy        = TSDestroy_Theta;
1246:   ts->ops->view           = TSView_Theta;
1247:   ts->ops->setup          = TSSetUp_Theta;
1248:   ts->ops->adjointsetup   = TSAdjointSetUp_Theta;
1249:   ts->ops->adjointreset   = TSAdjointReset_Theta;
1250:   ts->ops->step           = TSStep_Theta;
1251:   ts->ops->interpolate    = TSInterpolate_Theta;
1252:   ts->ops->evaluatewlte   = TSEvaluateWLTE_Theta;
1253:   ts->ops->rollback       = TSRollBack_Theta;
1254:   ts->ops->resizeregister = TSResizeRegister_Theta;
1255:   ts->ops->setfromoptions = TSSetFromOptions_Theta;
1256:   ts->ops->snesfunction   = SNESTSFormFunction_Theta;
1257:   ts->ops->snesjacobian   = SNESTSFormJacobian_Theta;
1258: #if defined(PETSC_HAVE_COMPLEX)
1259:   ts->ops->linearstability = TSComputeLinearStability_Theta;
1260: #endif
1261:   ts->ops->getstages       = TSGetStages_Theta;
1262:   ts->ops->adjointstep     = TSAdjointStep_Theta;
1263:   ts->ops->adjointintegral = TSAdjointCostIntegral_Theta;
1264:   ts->ops->forwardintegral = TSForwardCostIntegral_Theta;
1265:   ts->default_adapt_type   = TSADAPTNONE;

1267:   ts->ops->forwardsetup     = TSForwardSetUp_Theta;
1268:   ts->ops->forwardreset     = TSForwardReset_Theta;
1269:   ts->ops->forwardstep      = TSForwardStep_Theta;
1270:   ts->ops->forwardgetstages = TSForwardGetStages_Theta;

1272:   ts->usessnes = PETSC_TRUE;

1274:   PetscCall(PetscNew(&th));
1275:   ts->data = (void *)th;

1277:   th->VecsDeltaLam   = NULL;
1278:   th->VecsDeltaMu    = NULL;
1279:   th->VecsSensiTemp  = NULL;
1280:   th->VecsSensi2Temp = NULL;

1282:   th->extrapolate = PETSC_FALSE;
1283:   th->Theta       = 0.5;
1284:   th->order       = 2;
1285:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaGetTheta_C", TSThetaGetTheta_Theta));
1286:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaSetTheta_C", TSThetaSetTheta_Theta));
1287:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaGetEndpoint_C", TSThetaGetEndpoint_Theta));
1288:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaSetEndpoint_C", TSThetaSetEndpoint_Theta));
1289:   PetscFunctionReturn(PETSC_SUCCESS);
1290: }

1292: /*@
1293:   TSThetaGetTheta - Get the abscissa of the stage in (0,1] for `TSTHETA`

1295:   Not Collective

1297:   Input Parameter:
1298: . ts - timestepping context

1300:   Output Parameter:
1301: . theta - stage abscissa

1303:   Level: advanced

1305:   Note:
1306:   Use of this function is normally only required to hack `TSTHETA` to use a modified integration scheme.

1308: .seealso: [](ch_ts), `TSThetaSetTheta()`, `TSTHETA`
1309: @*/
1310: PetscErrorCode TSThetaGetTheta(TS ts, PetscReal *theta)
1311: {
1312:   PetscFunctionBegin;
1314:   PetscAssertPointer(theta, 2);
1315:   PetscUseMethod(ts, "TSThetaGetTheta_C", (TS, PetscReal *), (ts, theta));
1316:   PetscFunctionReturn(PETSC_SUCCESS);
1317: }

1319: /*@
1320:   TSThetaSetTheta - Set the abscissa of the stage in (0,1]  for `TSTHETA`

1322:   Not Collective

1324:   Input Parameters:
1325: + ts    - timestepping context
1326: - theta - stage abscissa

1328:   Options Database Key:
1329: . -ts_theta_theta <theta> - set theta

1331:   Level: intermediate

1333: .seealso: [](ch_ts), `TSThetaGetTheta()`, `TSTHETA`, `TSCN`
1334: @*/
1335: PetscErrorCode TSThetaSetTheta(TS ts, PetscReal theta)
1336: {
1337:   PetscFunctionBegin;
1339:   PetscTryMethod(ts, "TSThetaSetTheta_C", (TS, PetscReal), (ts, theta));
1340:   PetscFunctionReturn(PETSC_SUCCESS);
1341: }

1343: /*@
1344:   TSThetaGetEndpoint - Gets whether to use the endpoint variant of the method (e.g. trapezoid/Crank-Nicolson instead of midpoint rule) for `TSTHETA`

1346:   Not Collective

1348:   Input Parameter:
1349: . ts - timestepping context

1351:   Output Parameter:
1352: . endpoint - `PETSC_TRUE` when using the endpoint variant

1354:   Level: advanced

1356: .seealso: [](ch_ts), `TSThetaSetEndpoint()`, `TSTHETA`, `TSCN`
1357: @*/
1358: PetscErrorCode TSThetaGetEndpoint(TS ts, PetscBool *endpoint)
1359: {
1360:   PetscFunctionBegin;
1362:   PetscAssertPointer(endpoint, 2);
1363:   PetscUseMethod(ts, "TSThetaGetEndpoint_C", (TS, PetscBool *), (ts, endpoint));
1364:   PetscFunctionReturn(PETSC_SUCCESS);
1365: }

1367: /*@
1368:   TSThetaSetEndpoint - Sets whether to use the endpoint variant of the method (e.g. trapezoid/Crank-Nicolson instead of midpoint rule) for `TSTHETA`

1370:   Not Collective

1372:   Input Parameters:
1373: + ts  - timestepping context
1374: - flg - `PETSC_TRUE` to use the endpoint variant

1376:   Options Database Key:
1377: . -ts_theta_endpoint <flg> - use the endpoint variant

1379:   Level: intermediate

1381: .seealso: [](ch_ts), `TSTHETA`, `TSCN`
1382: @*/
1383: PetscErrorCode TSThetaSetEndpoint(TS ts, PetscBool flg)
1384: {
1385:   PetscFunctionBegin;
1387:   PetscTryMethod(ts, "TSThetaSetEndpoint_C", (TS, PetscBool), (ts, flg));
1388:   PetscFunctionReturn(PETSC_SUCCESS);
1389: }

1391: /*
1392:  * TSBEULER and TSCN are straightforward specializations of TSTHETA.
1393:  * The creation functions for these specializations are below.
1394:  */

1396: static PetscErrorCode TSSetUp_BEuler(TS ts)
1397: {
1398:   TS_Theta *th = (TS_Theta *)ts->data;

1400:   PetscFunctionBegin;
1401:   PetscCheck(th->Theta == 1.0, PetscObjectComm((PetscObject)ts), PETSC_ERR_OPT_OVERWRITE, "Can not change the default value (1) of theta when using backward Euler");
1402:   PetscCheck(!th->endpoint, PetscObjectComm((PetscObject)ts), PETSC_ERR_OPT_OVERWRITE, "Can not change to the endpoint form of the Theta methods when using backward Euler");
1403:   PetscCall(TSSetUp_Theta(ts));
1404:   PetscFunctionReturn(PETSC_SUCCESS);
1405: }

1407: static PetscErrorCode TSView_BEuler(TS ts, PetscViewer viewer)
1408: {
1409:   PetscFunctionBegin;
1410:   PetscFunctionReturn(PETSC_SUCCESS);
1411: }

1413: /*MC
1414:       TSBEULER - ODE solver using the implicit backward Euler method

1416:   Level: beginner

1418:   Note:
1419:   `TSBEULER` is equivalent to `TSTHETA` with Theta=1.0 or `-ts_type theta -ts_theta_theta 1.0`

1421: .seealso: [](ch_ts), `TSCreate()`, `TS`, `TSSetType()`, `TSEULER`, `TSCN`, `TSTHETA`
1422: M*/
1423: PETSC_EXTERN PetscErrorCode TSCreate_BEuler(TS ts)
1424: {
1425:   PetscFunctionBegin;
1426:   PetscCall(TSCreate_Theta(ts));
1427:   PetscCall(TSThetaSetTheta(ts, 1.0));
1428:   PetscCall(TSThetaSetEndpoint(ts, PETSC_FALSE));
1429:   ts->ops->setup = TSSetUp_BEuler;
1430:   ts->ops->view  = TSView_BEuler;
1431:   PetscFunctionReturn(PETSC_SUCCESS);
1432: }

1434: static PetscErrorCode TSSetUp_CN(TS ts)
1435: {
1436:   TS_Theta *th = (TS_Theta *)ts->data;

1438:   PetscFunctionBegin;
1439:   PetscCheck(th->Theta == 0.5, PetscObjectComm((PetscObject)ts), PETSC_ERR_OPT_OVERWRITE, "Can not change the default value (0.5) of theta when using Crank-Nicolson");
1440:   PetscCheck(th->endpoint, PetscObjectComm((PetscObject)ts), PETSC_ERR_OPT_OVERWRITE, "Can not change to the midpoint form of the Theta methods when using Crank-Nicolson");
1441:   PetscCall(TSSetUp_Theta(ts));
1442:   PetscFunctionReturn(PETSC_SUCCESS);
1443: }

1445: static PetscErrorCode TSView_CN(TS ts, PetscViewer viewer)
1446: {
1447:   PetscFunctionBegin;
1448:   PetscFunctionReturn(PETSC_SUCCESS);
1449: }

1451: /*MC
1452:       TSCN - ODE solver using the implicit Crank-Nicolson method.

1454:   Level: beginner

1456:   Notes:
1457:   `TSCN` is equivalent to `TSTHETA` with Theta=0.5 and the "endpoint" option set. I.e.
1458: .vb
1459:   -ts_type theta
1460:   -ts_theta_theta 0.5
1461:   -ts_theta_endpoint
1462: .ve

1464: .seealso: [](ch_ts), `TSCreate()`, `TS`, `TSSetType()`, `TSBEULER`, `TSTHETA`, `TSType`,
1465: M*/
1466: PETSC_EXTERN PetscErrorCode TSCreate_CN(TS ts)
1467: {
1468:   PetscFunctionBegin;
1469:   PetscCall(TSCreate_Theta(ts));
1470:   PetscCall(TSThetaSetTheta(ts, 0.5));
1471:   PetscCall(TSThetaSetEndpoint(ts, PETSC_TRUE));
1472:   ts->ops->setup = TSSetUp_CN;
1473:   ts->ops->view  = TSView_CN;
1474:   PetscFunctionReturn(PETSC_SUCCESS);
1475: }