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: /* We need to transfer X0 which will be copied into sol_prev */
183: static PetscErrorCode TSResizeRegister_Theta(TS ts, PetscBool reg)
184: {
185:   TS_Theta  *th     = (TS_Theta *)ts->data;
186:   const char name[] = "ts:theta:X0";

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

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

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

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

229:     th->status = TS_STEP_PENDING;
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));
234:       PetscCall(VecAXPY(ts->vec_sol, ts->time_step, th->Xdot));
235:     }
236:     PetscCall(TSAdaptChoose(ts->adapt, ts, ts->time_step, NULL, &next_time_step, &accept));
237:     th->status = accept ? TS_STEP_COMPLETE : TS_STEP_INCOMPLETE;
238:     if (!accept) {
239:       PetscCall(VecCopy(th->X0, ts->vec_sol));
240:       ts->time_step = next_time_step;
241:       goto reject_step;
242:     }

244:     if (ts->forward_solve || ts->costintegralfwd) { /* Save the info for the later use in cost integral evaluation */
245:       th->ptime0     = ts->ptime;
246:       th->time_step0 = ts->time_step;
247:     }
248:     ts->ptime += ts->time_step;
249:     ts->time_step = next_time_step;
250:     break;

252:   reject_step:
253:     ts->reject++;
254:     accept = PETSC_FALSE;
255:     if (!ts->reason && ++rejections > ts->max_reject && ts->max_reject >= 0) {
256:       ts->reason = TS_DIVERGED_STEP_REJECTED;
257:       PetscCall(PetscInfo(ts, "Step=%" PetscInt_FMT ", step rejections %" PetscInt_FMT " greater than current TS allowed, stopping solve\n", ts->steps, rejections));
258:     }
259:   }
260:   PetscFunctionReturn(PETSC_SUCCESS);
261: }

263: static PetscErrorCode TSAdjointStepBEuler_Private(TS ts)
264: {
265:   TS_Theta      *th           = (TS_Theta *)ts->data;
266:   TS             quadts       = ts->quadraturets;
267:   Vec           *VecsDeltaLam = th->VecsDeltaLam, *VecsDeltaMu = th->VecsDeltaMu, *VecsSensiTemp = th->VecsSensiTemp;
268:   Vec           *VecsDeltaLam2 = th->VecsDeltaLam2, *VecsDeltaMu2 = th->VecsDeltaMu2, *VecsSensi2Temp = th->VecsSensi2Temp;
269:   PetscInt       nadj;
270:   Mat            J, Jpre, quadJ = NULL, quadJp = NULL;
271:   KSP            ksp;
272:   PetscScalar   *xarr;
273:   TSEquationType eqtype;
274:   PetscBool      isexplicitode = PETSC_FALSE;
275:   PetscReal      adjoint_time_step;

277:   PetscFunctionBegin;
278:   PetscCall(TSGetEquationType(ts, &eqtype));
279:   if (eqtype == TS_EQ_ODE_EXPLICIT) {
280:     isexplicitode = PETSC_TRUE;
281:     VecsDeltaLam  = ts->vecs_sensi;
282:     VecsDeltaLam2 = ts->vecs_sensi2;
283:   }
284:   th->status = TS_STEP_INCOMPLETE;
285:   PetscCall(SNESGetKSP(ts->snes, &ksp));
286:   PetscCall(TSGetIJacobian(ts, &J, &Jpre, NULL, NULL));
287:   if (quadts) {
288:     PetscCall(TSGetRHSJacobian(quadts, &quadJ, NULL, NULL, NULL));
289:     PetscCall(TSGetRHSJacobianP(quadts, &quadJp, NULL, NULL));
290:   }

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

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

298:   for (nadj = 0; nadj < ts->numcost; nadj++) {
299:     PetscCall(VecCopy(ts->vecs_sensi[nadj], VecsSensiTemp[nadj]));
300:     PetscCall(VecScale(VecsSensiTemp[nadj], 1. / adjoint_time_step)); /* lambda_{n+1}/h */
301:     if (quadJ) {
302:       PetscCall(MatDenseGetColumn(quadJ, nadj, &xarr));
303:       PetscCall(VecPlaceArray(ts->vec_drdu_col, xarr));
304:       PetscCall(VecAXPY(VecsSensiTemp[nadj], 1., ts->vec_drdu_col));
305:       PetscCall(VecResetArray(ts->vec_drdu_col));
306:       PetscCall(MatDenseRestoreColumn(quadJ, &xarr));
307:     }
308:   }

310:   /* Build LHS for first-order adjoint */
311:   th->shift = 1. / adjoint_time_step;
312:   PetscCall(TSComputeSNESJacobian(ts, ts->vec_sol, J, Jpre));
313:   PetscCall(KSPSetOperators(ksp, J, Jpre));

315:   /* Solve stage equation LHS*lambda_s = RHS for first-order adjoint */
316:   for (nadj = 0; nadj < ts->numcost; nadj++) {
317:     KSPConvergedReason kspreason;
318:     PetscCall(KSPSolveTranspose(ksp, VecsSensiTemp[nadj], VecsDeltaLam[nadj]));
319:     PetscCall(KSPGetConvergedReason(ksp, &kspreason));
320:     if (kspreason < 0) {
321:       ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
322:       PetscCall(PetscInfo(ts, "Step=%" PetscInt_FMT ", %" PetscInt_FMT "th cost function, transposed linear solve fails, stopping 1st-order adjoint solve\n", ts->steps, nadj));
323:     }
324:   }

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

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

381:     for (nadj = 0; nadj < ts->numcost; nadj++) {
382:       PetscCall(MatMultTranspose(ts->Jacp, VecsDeltaLam[nadj], VecsDeltaMu[nadj]));
383:       PetscCall(VecAXPY(ts->vecs_sensip[nadj], -adjoint_time_step, VecsDeltaMu[nadj]));
384:       if (quadJp) {
385:         PetscCall(MatDenseGetColumn(quadJp, nadj, &xarr));
386:         PetscCall(VecPlaceArray(ts->vec_drdp_col, xarr));
387:         PetscCall(VecAXPY(ts->vecs_sensip[nadj], adjoint_time_step, ts->vec_drdp_col));
388:         PetscCall(VecResetArray(ts->vec_drdp_col));
389:         PetscCall(MatDenseRestoreColumn(quadJp, &xarr));
390:       }
391:       if (ts->vecs_sensi2p) {
392:         PetscCall(MatMultTranspose(ts->Jacp, VecsDeltaLam2[nadj], VecsDeltaMu2[nadj]));
393:         PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step, VecsDeltaMu2[nadj]));
394:         if (ts->vecs_fpu) PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step, ts->vecs_fpu[nadj]));
395:         if (ts->vecs_fpp) PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step, ts->vecs_fpp[nadj]));
396:       }
397:     }
398:   }

400:   if (ts->vecs_sensi2) {
401:     PetscCall(VecResetArray(ts->vec_sensip_col));
402:     PetscCall(MatDenseRestoreColumn(ts->mat_sensip, &xarr));
403:   }
404:   th->status = TS_STEP_COMPLETE;
405:   PetscFunctionReturn(PETSC_SUCCESS);
406: }

408: static PetscErrorCode TSAdjointStep_Theta(TS ts)
409: {
410:   TS_Theta    *th           = (TS_Theta *)ts->data;
411:   TS           quadts       = ts->quadraturets;
412:   Vec         *VecsDeltaLam = th->VecsDeltaLam, *VecsDeltaMu = th->VecsDeltaMu, *VecsSensiTemp = th->VecsSensiTemp;
413:   Vec         *VecsDeltaLam2 = th->VecsDeltaLam2, *VecsDeltaMu2 = th->VecsDeltaMu2, *VecsSensi2Temp = th->VecsSensi2Temp;
414:   PetscInt     nadj;
415:   Mat          J, Jpre, quadJ = NULL, quadJp = NULL;
416:   KSP          ksp;
417:   PetscScalar *xarr;
418:   PetscReal    adjoint_time_step;
419:   PetscReal    adjoint_ptime; /* end time of the adjoint time step (ts->ptime is the start time, usually ts->ptime is larger than adjoint_ptime) */

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

438:   if (!th->endpoint) {
439:     /* recover th->X0 using vec_sol and the stage value th->X */
440:     PetscCall(VecAXPBYPCZ(th->X0, 1.0 / (1.0 - th->Theta), th->Theta / (th->Theta - 1.0), 0, th->X, ts->vec_sol));
441:   }

443:   /* Build RHS for first-order adjoint */
444:   /* Cost function has an integral term */
445:   if (quadts) {
446:     if (th->endpoint) {
447:       PetscCall(TSComputeRHSJacobian(quadts, th->stage_time, ts->vec_sol, quadJ, NULL));
448:     } else {
449:       PetscCall(TSComputeRHSJacobian(quadts, th->stage_time, th->X, quadJ, NULL));
450:     }
451:   }

453:   for (nadj = 0; nadj < ts->numcost; nadj++) {
454:     PetscCall(VecCopy(ts->vecs_sensi[nadj], VecsSensiTemp[nadj]));
455:     PetscCall(VecScale(VecsSensiTemp[nadj], 1. / (th->Theta * adjoint_time_step)));
456:     if (quadJ) {
457:       PetscCall(MatDenseGetColumn(quadJ, nadj, &xarr));
458:       PetscCall(VecPlaceArray(ts->vec_drdu_col, xarr));
459:       PetscCall(VecAXPY(VecsSensiTemp[nadj], 1., ts->vec_drdu_col));
460:       PetscCall(VecResetArray(ts->vec_drdu_col));
461:       PetscCall(MatDenseRestoreColumn(quadJ, &xarr));
462:     }
463:   }

465:   /* Build LHS for first-order adjoint */
466:   th->shift = 1. / (th->Theta * adjoint_time_step);
467:   if (th->endpoint) {
468:     PetscCall(TSComputeSNESJacobian(ts, ts->vec_sol, J, Jpre));
469:   } else {
470:     PetscCall(TSComputeSNESJacobian(ts, th->X, J, Jpre));
471:   }
472:   PetscCall(KSPSetOperators(ksp, J, Jpre));

474:   /* Solve stage equation LHS*lambda_s = RHS for first-order adjoint */
475:   for (nadj = 0; nadj < ts->numcost; nadj++) {
476:     KSPConvergedReason kspreason;
477:     PetscCall(KSPSolveTranspose(ksp, VecsSensiTemp[nadj], VecsDeltaLam[nadj]));
478:     PetscCall(KSPGetConvergedReason(ksp, &kspreason));
479:     if (kspreason < 0) {
480:       ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
481:       PetscCall(PetscInfo(ts, "Step=%" PetscInt_FMT ", %" PetscInt_FMT "th cost function, transposed linear solve fails, stopping 1st-order adjoint solve\n", ts->steps, nadj));
482:     }
483:   }

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

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

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

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

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

583:         /* lambda_s^T F_PP w_2 */
584:         PetscCall(TSComputeIHessianProductFunctionPP(ts, th->stage_time, ts->vec_sol, VecsDeltaLam, ts->vec_dir, ts->vecs_fpp));
585:         for (nadj = 0; nadj < ts->numcost; nadj++) {
586:           /* 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)  */
587:           PetscCall(MatMultTranspose(ts->Jacp, VecsDeltaLam2[nadj], VecsDeltaMu2[nadj]));
588:           PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step * th->Theta, VecsDeltaMu2[nadj]));
589:           if (ts->vecs_fpu) PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step * th->Theta, ts->vecs_fpu[nadj]));
590:           if (ts->vecs_fpp) PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step * th->Theta, ts->vecs_fpp[nadj]));
591:         }
592:       }

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

663:   th->status = TS_STEP_COMPLETE;
664:   PetscFunctionReturn(PETSC_SUCCESS);
665: }

667: static PetscErrorCode TSInterpolate_Theta(TS ts, PetscReal t, Vec X)
668: {
669:   TS_Theta *th = (TS_Theta *)ts->data;
670:   PetscReal dt = t - ts->ptime;

672:   PetscFunctionBegin;
673:   PetscCall(VecCopy(ts->vec_sol, th->X));
674:   if (th->endpoint) dt *= th->Theta;
675:   PetscCall(VecWAXPY(X, dt, th->Xdot, th->X));
676:   PetscFunctionReturn(PETSC_SUCCESS);
677: }

679: static PetscErrorCode TSEvaluateWLTE_Theta(TS ts, NormType wnormtype, PetscInt *order, PetscReal *wlte)
680: {
681:   TS_Theta *th = (TS_Theta *)ts->data;
682:   Vec       X  = ts->vec_sol;      /* X = solution */
683:   Vec       Y  = th->vec_lte_work; /* Y = X + LTE  */
684:   PetscReal wltea, wlter;

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

716: static PetscErrorCode TSRollBack_Theta(TS ts)
717: {
718:   TS_Theta *th     = (TS_Theta *)ts->data;
719:   TS        quadts = ts->quadraturets;

721:   PetscFunctionBegin;
722:   PetscCall(VecCopy(th->X0, ts->vec_sol));
723:   if (quadts && ts->costintegralfwd) PetscCall(VecCopy(th->VecCostIntegral0, quadts->vec_sol));
724:   th->status = TS_STEP_INCOMPLETE;
725:   if (ts->mat_sensip) PetscCall(MatCopy(th->MatFwdSensip0, ts->mat_sensip, SAME_NONZERO_PATTERN));
726:   if (quadts && quadts->mat_sensip) PetscCall(MatCopy(th->MatIntegralSensip0, quadts->mat_sensip, SAME_NONZERO_PATTERN));
727:   PetscFunctionReturn(PETSC_SUCCESS);
728: }

730: static PetscErrorCode TSForwardStep_Theta(TS ts)
731: {
732:   TS_Theta    *th                   = (TS_Theta *)ts->data;
733:   TS           quadts               = ts->quadraturets;
734:   Mat          MatDeltaFwdSensip    = th->MatDeltaFwdSensip;
735:   Vec          VecDeltaFwdSensipCol = th->VecDeltaFwdSensipCol;
736:   PetscInt     ntlm;
737:   KSP          ksp;
738:   Mat          J, Jpre, quadJ = NULL, quadJp = NULL;
739:   PetscScalar *barr, *xarr;
740:   PetscReal    previous_shift;

742:   PetscFunctionBegin;
743:   previous_shift = th->shift;
744:   PetscCall(MatCopy(ts->mat_sensip, th->MatFwdSensip0, SAME_NONZERO_PATTERN));

746:   if (quadts && quadts->mat_sensip) PetscCall(MatCopy(quadts->mat_sensip, th->MatIntegralSensip0, SAME_NONZERO_PATTERN));
747:   PetscCall(SNESGetKSP(ts->snes, &ksp));
748:   PetscCall(TSGetIJacobian(ts, &J, &Jpre, NULL, NULL));
749:   if (quadts) {
750:     PetscCall(TSGetRHSJacobian(quadts, &quadJ, NULL, NULL, NULL));
751:     PetscCall(TSGetRHSJacobianP(quadts, &quadJp, NULL, NULL));
752:   }

754:   /* Build RHS */
755:   if (th->endpoint) { /* 2-stage method*/
756:     th->shift = 1. / ((th->Theta - 1.) * th->time_step0);
757:     PetscCall(TSComputeIJacobian(ts, th->ptime0, th->X0, th->Xdot, th->shift, J, Jpre, PETSC_FALSE));
758:     PetscCall(MatMatMult(J, ts->mat_sensip, MAT_REUSE_MATRIX, PETSC_DEFAULT, &MatDeltaFwdSensip));
759:     PetscCall(MatScale(MatDeltaFwdSensip, (th->Theta - 1.) / th->Theta));

761:     /* Add the f_p forcing terms */
762:     if (ts->Jacp) {
763:       PetscCall(VecZeroEntries(th->Xdot));
764:       PetscCall(TSComputeIJacobianP(ts, th->ptime0, th->X0, th->Xdot, th->shift, ts->Jacp, PETSC_FALSE));
765:       PetscCall(MatAXPY(MatDeltaFwdSensip, (th->Theta - 1.) / th->Theta, ts->Jacp, SUBSET_NONZERO_PATTERN));
766:       th->shift = previous_shift;
767:       PetscCall(VecAXPBYPCZ(th->Xdot, -th->shift, th->shift, 0, th->X0, ts->vec_sol));
768:       PetscCall(TSComputeIJacobianP(ts, th->stage_time, ts->vec_sol, th->Xdot, th->shift, ts->Jacp, PETSC_FALSE));
769:       PetscCall(MatAXPY(MatDeltaFwdSensip, -1., ts->Jacp, SUBSET_NONZERO_PATTERN));
770:     }
771:   } else { /* 1-stage method */
772:     th->shift = 0.0;
773:     PetscCall(TSComputeIJacobian(ts, th->stage_time, th->X, th->Xdot, th->shift, J, Jpre, PETSC_FALSE));
774:     PetscCall(MatMatMult(J, ts->mat_sensip, MAT_REUSE_MATRIX, PETSC_DEFAULT, &MatDeltaFwdSensip));
775:     PetscCall(MatScale(MatDeltaFwdSensip, -1.));

777:     /* Add the f_p forcing terms */
778:     if (ts->Jacp) {
779:       th->shift = previous_shift;
780:       PetscCall(VecAXPBYPCZ(th->Xdot, -th->shift, th->shift, 0, th->X0, th->X));
781:       PetscCall(TSComputeIJacobianP(ts, th->stage_time, th->X, th->Xdot, th->shift, ts->Jacp, PETSC_FALSE));
782:       PetscCall(MatAXPY(MatDeltaFwdSensip, -1., ts->Jacp, SUBSET_NONZERO_PATTERN));
783:     }
784:   }

786:   /* Build LHS */
787:   th->shift = previous_shift; /* recover the previous shift used in TSStep_Theta() */
788:   if (th->endpoint) {
789:     PetscCall(TSComputeIJacobian(ts, th->stage_time, ts->vec_sol, th->Xdot, th->shift, J, Jpre, PETSC_FALSE));
790:   } else {
791:     PetscCall(TSComputeIJacobian(ts, th->stage_time, th->X, th->Xdot, th->shift, J, Jpre, PETSC_FALSE));
792:   }
793:   PetscCall(KSPSetOperators(ksp, J, Jpre));

795:   /*
796:     Evaluate the first stage of integral gradients with the 2-stage method:
797:     drdu|t_n*S(t_n) + drdp|t_n
798:     This is done before the linear solve because the sensitivity variable S(t_n) will be propagated to S(t_{n+1})
799:   */
800:   if (th->endpoint) { /* 2-stage method only */
801:     if (quadts && quadts->mat_sensip) {
802:       PetscCall(TSComputeRHSJacobian(quadts, th->ptime0, th->X0, quadJ, NULL));
803:       PetscCall(TSComputeRHSJacobianP(quadts, th->ptime0, th->X0, quadJp));
804:       PetscCall(MatTransposeMatMult(ts->mat_sensip, quadJ, MAT_REUSE_MATRIX, PETSC_DEFAULT, &th->MatIntegralSensipTemp));
805:       PetscCall(MatAXPY(th->MatIntegralSensipTemp, 1, quadJp, SAME_NONZERO_PATTERN));
806:       PetscCall(MatAXPY(quadts->mat_sensip, th->time_step0 * (1. - th->Theta), th->MatIntegralSensipTemp, SAME_NONZERO_PATTERN));
807:     }
808:   }

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

833:   /*
834:     Evaluate the second stage of integral gradients with the 2-stage method:
835:     drdu|t_{n+1}*S(t_{n+1}) + drdp|t_{n+1}
836:   */
837:   if (quadts && quadts->mat_sensip) {
838:     if (!th->endpoint) {
839:       PetscCall(MatAXPY(ts->mat_sensip, 1, MatDeltaFwdSensip, SAME_NONZERO_PATTERN)); /* stage sensitivity */
840:       PetscCall(TSComputeRHSJacobian(quadts, th->stage_time, th->X, quadJ, NULL));
841:       PetscCall(TSComputeRHSJacobianP(quadts, th->stage_time, th->X, quadJp));
842:       PetscCall(MatTransposeMatMult(ts->mat_sensip, quadJ, MAT_REUSE_MATRIX, PETSC_DEFAULT, &th->MatIntegralSensipTemp));
843:       PetscCall(MatAXPY(th->MatIntegralSensipTemp, 1, quadJp, SAME_NONZERO_PATTERN));
844:       PetscCall(MatAXPY(quadts->mat_sensip, th->time_step0, th->MatIntegralSensipTemp, SAME_NONZERO_PATTERN));
845:       PetscCall(MatAXPY(ts->mat_sensip, (1. - th->Theta) / th->Theta, MatDeltaFwdSensip, SAME_NONZERO_PATTERN));
846:     } else {
847:       PetscCall(TSComputeRHSJacobian(quadts, th->stage_time, ts->vec_sol, quadJ, NULL));
848:       PetscCall(TSComputeRHSJacobianP(quadts, th->stage_time, ts->vec_sol, quadJp));
849:       PetscCall(MatTransposeMatMult(ts->mat_sensip, quadJ, MAT_REUSE_MATRIX, PETSC_DEFAULT, &th->MatIntegralSensipTemp));
850:       PetscCall(MatAXPY(th->MatIntegralSensipTemp, 1, quadJp, SAME_NONZERO_PATTERN));
851:       PetscCall(MatAXPY(quadts->mat_sensip, th->time_step0 * th->Theta, th->MatIntegralSensipTemp, SAME_NONZERO_PATTERN));
852:     }
853:   } else {
854:     if (!th->endpoint) PetscCall(MatAXPY(ts->mat_sensip, 1. / th->Theta, MatDeltaFwdSensip, SAME_NONZERO_PATTERN));
855:   }
856:   PetscFunctionReturn(PETSC_SUCCESS);
857: }

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

863:   PetscFunctionBegin;
864:   if (ns) {
865:     if (!th->endpoint && th->Theta != 1.0) *ns = 1; /* midpoint form */
866:     else *ns = 2;                                   /* endpoint form */
867:   }
868:   if (stagesensip) {
869:     if (!th->endpoint && th->Theta != 1.0) {
870:       th->MatFwdStages[0] = th->MatDeltaFwdSensip;
871:     } else {
872:       th->MatFwdStages[0] = th->MatFwdSensip0;
873:       th->MatFwdStages[1] = ts->mat_sensip; /* stiffly accurate */
874:     }
875:     *stagesensip = th->MatFwdStages;
876:   }
877:   PetscFunctionReturn(PETSC_SUCCESS);
878: }

880: /*------------------------------------------------------------*/
881: static PetscErrorCode TSReset_Theta(TS ts)
882: {
883:   TS_Theta *th = (TS_Theta *)ts->data;

885:   PetscFunctionBegin;
886:   PetscCall(VecDestroy(&th->X));
887:   PetscCall(VecDestroy(&th->Xdot));
888:   PetscCall(VecDestroy(&th->X0));
889:   PetscCall(VecDestroy(&th->affine));

891:   PetscCall(VecDestroy(&th->vec_sol_prev));
892:   PetscCall(VecDestroy(&th->vec_lte_work));

894:   PetscCall(VecDestroy(&th->VecCostIntegral0));
895:   PetscFunctionReturn(PETSC_SUCCESS);
896: }

898: static PetscErrorCode TSAdjointReset_Theta(TS ts)
899: {
900:   TS_Theta *th = (TS_Theta *)ts->data;

902:   PetscFunctionBegin;
903:   PetscCall(VecDestroyVecs(ts->numcost, &th->VecsDeltaLam));
904:   PetscCall(VecDestroyVecs(ts->numcost, &th->VecsDeltaMu));
905:   PetscCall(VecDestroyVecs(ts->numcost, &th->VecsDeltaLam2));
906:   PetscCall(VecDestroyVecs(ts->numcost, &th->VecsDeltaMu2));
907:   PetscCall(VecDestroyVecs(ts->numcost, &th->VecsSensiTemp));
908:   PetscCall(VecDestroyVecs(ts->numcost, &th->VecsSensi2Temp));
909:   PetscFunctionReturn(PETSC_SUCCESS);
910: }

912: static PetscErrorCode TSDestroy_Theta(TS ts)
913: {
914:   PetscFunctionBegin;
915:   PetscCall(TSReset_Theta(ts));
916:   if (ts->dm) {
917:     PetscCall(DMCoarsenHookRemove(ts->dm, DMCoarsenHook_TSTheta, DMRestrictHook_TSTheta, ts));
918:     PetscCall(DMSubDomainHookRemove(ts->dm, DMSubDomainHook_TSTheta, DMSubDomainRestrictHook_TSTheta, ts));
919:   }
920:   PetscCall(PetscFree(ts->data));
921:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaGetTheta_C", NULL));
922:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaSetTheta_C", NULL));
923:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaGetEndpoint_C", NULL));
924:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaSetEndpoint_C", NULL));
925:   PetscFunctionReturn(PETSC_SUCCESS);
926: }

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

932:   Note that U here is the stage argument. This means that U = U_{n+1} only if endpoint = true,
933:   otherwise U = theta U_{n+1} + (1 - theta) U0, which for the case of implicit midpoint is
934:   U = (U_{n+1} + U0)/2
935: */
936: static PetscErrorCode SNESTSFormFunction_Theta(SNES snes, Vec x, Vec y, TS ts)
937: {
938:   TS_Theta *th = (TS_Theta *)ts->data;
939:   Vec       X0, Xdot;
940:   DM        dm, dmsave;
941:   PetscReal shift = th->shift;

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

961: static PetscErrorCode SNESTSFormJacobian_Theta(SNES snes, Vec x, Mat A, Mat B, TS ts)
962: {
963:   TS_Theta *th = (TS_Theta *)ts->data;
964:   Vec       Xdot;
965:   DM        dm, dmsave;
966:   PetscReal shift = th->shift;

968:   PetscFunctionBegin;
969:   PetscCall(SNESGetDM(snes, &dm));
970:   /* Xdot has already been computed in SNESTSFormFunction_Theta (SNES guarantees this) */
971:   PetscCall(TSThetaGetX0AndXdot(ts, dm, NULL, &Xdot));

973:   dmsave = ts->dm;
974:   ts->dm = dm;
975:   PetscCall(TSComputeIJacobian(ts, th->stage_time, x, Xdot, shift, A, B, PETSC_FALSE));
976:   ts->dm = dmsave;
977:   PetscCall(TSThetaRestoreX0AndXdot(ts, dm, NULL, &Xdot));
978:   PetscFunctionReturn(PETSC_SUCCESS);
979: }

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

986:   PetscFunctionBegin;
987:   /* combine sensitivities to parameters and sensitivities to initial values into one array */
988:   th->num_tlm = ts->num_parameters;
989:   PetscCall(MatDuplicate(ts->mat_sensip, MAT_DO_NOT_COPY_VALUES, &th->MatDeltaFwdSensip));
990:   if (quadts && quadts->mat_sensip) {
991:     PetscCall(MatDuplicate(quadts->mat_sensip, MAT_DO_NOT_COPY_VALUES, &th->MatIntegralSensipTemp));
992:     PetscCall(MatDuplicate(quadts->mat_sensip, MAT_DO_NOT_COPY_VALUES, &th->MatIntegralSensip0));
993:   }
994:   /* backup sensitivity results for roll-backs */
995:   PetscCall(MatDuplicate(ts->mat_sensip, MAT_DO_NOT_COPY_VALUES, &th->MatFwdSensip0));

997:   PetscCall(VecDuplicate(ts->vec_sol, &th->VecDeltaFwdSensipCol));
998:   PetscFunctionReturn(PETSC_SUCCESS);
999: }

1001: static PetscErrorCode TSForwardReset_Theta(TS ts)
1002: {
1003:   TS_Theta *th     = (TS_Theta *)ts->data;
1004:   TS        quadts = ts->quadraturets;

1006:   PetscFunctionBegin;
1007:   if (quadts && quadts->mat_sensip) {
1008:     PetscCall(MatDestroy(&th->MatIntegralSensipTemp));
1009:     PetscCall(MatDestroy(&th->MatIntegralSensip0));
1010:   }
1011:   PetscCall(VecDestroy(&th->VecDeltaFwdSensipCol));
1012:   PetscCall(MatDestroy(&th->MatDeltaFwdSensip));
1013:   PetscCall(MatDestroy(&th->MatFwdSensip0));
1014:   PetscFunctionReturn(PETSC_SUCCESS);
1015: }

1017: static PetscErrorCode TSSetUp_Theta(TS ts)
1018: {
1019:   TS_Theta *th     = (TS_Theta *)ts->data;
1020:   TS        quadts = ts->quadraturets;
1021:   PetscBool match;

1023:   PetscFunctionBegin;
1024:   if (!th->VecCostIntegral0 && quadts && ts->costintegralfwd) { /* back up cost integral */
1025:     PetscCall(VecDuplicate(quadts->vec_sol, &th->VecCostIntegral0));
1026:   }
1027:   if (!th->X) PetscCall(VecDuplicate(ts->vec_sol, &th->X));
1028:   if (!th->Xdot) PetscCall(VecDuplicate(ts->vec_sol, &th->Xdot));
1029:   if (!th->X0) PetscCall(VecDuplicate(ts->vec_sol, &th->X0));
1030:   if (th->endpoint) PetscCall(VecDuplicate(ts->vec_sol, &th->affine));

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

1035:   PetscCall(TSGetDM(ts, &ts->dm));
1036:   PetscCall(DMCoarsenHookAdd(ts->dm, DMCoarsenHook_TSTheta, DMRestrictHook_TSTheta, ts));
1037:   PetscCall(DMSubDomainHookAdd(ts->dm, DMSubDomainHook_TSTheta, DMSubDomainRestrictHook_TSTheta, ts));

1039:   PetscCall(TSGetAdapt(ts, &ts->adapt));
1040:   PetscCall(TSAdaptCandidatesClear(ts->adapt));
1041:   PetscCall(PetscObjectTypeCompare((PetscObject)ts->adapt, TSADAPTNONE, &match));
1042:   if (!match) {
1043:     PetscCall(VecDuplicate(ts->vec_sol, &th->vec_sol_prev));
1044:     PetscCall(VecDuplicate(ts->vec_sol, &th->vec_lte_work));
1045:   }
1046:   PetscCall(TSGetSNES(ts, &ts->snes));

1048:   ts->stifflyaccurate = (!th->endpoint && th->Theta != 1.0) ? PETSC_FALSE : PETSC_TRUE;
1049:   PetscFunctionReturn(PETSC_SUCCESS);
1050: }

1052: /*------------------------------------------------------------*/

1054: static PetscErrorCode TSAdjointSetUp_Theta(TS ts)
1055: {
1056:   TS_Theta *th = (TS_Theta *)ts->data;

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

1078: static PetscErrorCode TSSetFromOptions_Theta(TS ts, PetscOptionItems *PetscOptionsObject)
1079: {
1080:   TS_Theta *th = (TS_Theta *)ts->data;

1082:   PetscFunctionBegin;
1083:   PetscOptionsHeadBegin(PetscOptionsObject, "Theta ODE solver options");
1084:   {
1085:     PetscCall(PetscOptionsReal("-ts_theta_theta", "Location of stage (0<Theta<=1)", "TSThetaSetTheta", th->Theta, &th->Theta, NULL));
1086:     PetscCall(PetscOptionsBool("-ts_theta_endpoint", "Use the endpoint instead of midpoint form of the Theta method", "TSThetaSetEndpoint", th->endpoint, &th->endpoint, NULL));
1087:     PetscCall(PetscOptionsBool("-ts_theta_initial_guess_extrapolate", "Extrapolate stage initial guess from previous solution (sometimes unstable)", "TSThetaSetExtrapolate", th->extrapolate, &th->extrapolate, NULL));
1088:   }
1089:   PetscOptionsHeadEnd();
1090:   PetscFunctionReturn(PETSC_SUCCESS);
1091: }

1093: static PetscErrorCode TSView_Theta(TS ts, PetscViewer viewer)
1094: {
1095:   TS_Theta *th = (TS_Theta *)ts->data;
1096:   PetscBool iascii;

1098:   PetscFunctionBegin;
1099:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &iascii));
1100:   if (iascii) {
1101:     PetscCall(PetscViewerASCIIPrintf(viewer, "  Theta=%g\n", (double)th->Theta));
1102:     PetscCall(PetscViewerASCIIPrintf(viewer, "  Extrapolation=%s\n", th->extrapolate ? "yes" : "no"));
1103:   }
1104:   PetscFunctionReturn(PETSC_SUCCESS);
1105: }

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

1111:   PetscFunctionBegin;
1112:   *theta = th->Theta;
1113:   PetscFunctionReturn(PETSC_SUCCESS);
1114: }

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

1120:   PetscFunctionBegin;
1121:   PetscCheck(theta > 0 && theta <= 1, PetscObjectComm((PetscObject)ts), PETSC_ERR_ARG_OUTOFRANGE, "Theta %g not in range (0,1]", (double)theta);
1122:   th->Theta = theta;
1123:   th->order = (th->Theta == 0.5) ? 2 : 1;
1124:   PetscFunctionReturn(PETSC_SUCCESS);
1125: }

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

1131:   PetscFunctionBegin;
1132:   *endpoint = th->endpoint;
1133:   PetscFunctionReturn(PETSC_SUCCESS);
1134: }

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

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

1145: #if defined(PETSC_HAVE_COMPLEX)
1146: static PetscErrorCode TSComputeLinearStability_Theta(TS ts, PetscReal xr, PetscReal xi, PetscReal *yr, PetscReal *yi)
1147: {
1148:   PetscComplex    z   = xr + xi * PETSC_i, f;
1149:   TS_Theta       *th  = (TS_Theta *)ts->data;
1150:   const PetscReal one = 1.0;

1152:   PetscFunctionBegin;
1153:   f   = (one + (one - th->Theta) * z) / (one - th->Theta * z);
1154:   *yr = PetscRealPartComplex(f);
1155:   *yi = PetscImaginaryPartComplex(f);
1156:   PetscFunctionReturn(PETSC_SUCCESS);
1157: }
1158: #endif

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

1164:   PetscFunctionBegin;
1165:   if (ns) {
1166:     if (!th->endpoint && th->Theta != 1.0) *ns = 1; /* midpoint form */
1167:     else *ns = 2;                                   /* endpoint form */
1168:   }
1169:   if (Y) {
1170:     if (!th->endpoint && th->Theta != 1.0) {
1171:       th->Stages[0] = th->X;
1172:     } else {
1173:       th->Stages[0] = th->X0;
1174:       th->Stages[1] = ts->vec_sol; /* stiffly accurate */
1175:     }
1176:     *Y = th->Stages;
1177:   }
1178:   PetscFunctionReturn(PETSC_SUCCESS);
1179: }

1181: /* ------------------------------------------------------------ */
1182: /*MC
1183:       TSTHETA - DAE solver using the implicit Theta method

1185:    Level: beginner

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

1192:    Notes:
1193: .vb
1194:   -ts_type theta -ts_theta_theta 1.0 corresponds to backward Euler (TSBEULER)
1195:   -ts_type theta -ts_theta_theta 0.5 corresponds to the implicit midpoint rule
1196:   -ts_type theta -ts_theta_theta 0.5 -ts_theta_endpoint corresponds to Crank-Nicholson (TSCN)
1197: .ve

1199:    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.

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

1203: .vb
1204:   Theta | Theta
1205:   -------------
1206:         |  1
1207: .ve

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

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

1213: .vb
1214:   0 | 0         0
1215:   1 | 1-Theta   Theta
1216:   -------------------
1217:     | 1-Theta   Theta
1218: .ve

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

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

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

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

1228: .seealso: [](ch_ts), `TSCreate()`, `TS`, `TSSetType()`, `TSCN`, `TSBEULER`, `TSThetaSetTheta()`, `TSThetaSetEndpoint()`
1229: M*/
1230: PETSC_EXTERN PetscErrorCode TSCreate_Theta(TS ts)
1231: {
1232:   TS_Theta *th;

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

1259:   ts->ops->forwardsetup     = TSForwardSetUp_Theta;
1260:   ts->ops->forwardreset     = TSForwardReset_Theta;
1261:   ts->ops->forwardstep      = TSForwardStep_Theta;
1262:   ts->ops->forwardgetstages = TSForwardGetStages_Theta;

1264:   ts->usessnes = PETSC_TRUE;

1266:   PetscCall(PetscNew(&th));
1267:   ts->data = (void *)th;

1269:   th->VecsDeltaLam   = NULL;
1270:   th->VecsDeltaMu    = NULL;
1271:   th->VecsSensiTemp  = NULL;
1272:   th->VecsSensi2Temp = NULL;

1274:   th->extrapolate = PETSC_FALSE;
1275:   th->Theta       = 0.5;
1276:   th->order       = 2;
1277:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaGetTheta_C", TSThetaGetTheta_Theta));
1278:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaSetTheta_C", TSThetaSetTheta_Theta));
1279:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaGetEndpoint_C", TSThetaGetEndpoint_Theta));
1280:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaSetEndpoint_C", TSThetaSetEndpoint_Theta));
1281:   PetscFunctionReturn(PETSC_SUCCESS);
1282: }

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

1287:   Not Collective

1289:   Input Parameter:
1290: . ts - timestepping context

1292:   Output Parameter:
1293: . theta - stage abscissa

1295:   Level: advanced

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

1300: .seealso: [](ch_ts), `TSThetaSetTheta()`, `TSTHETA`
1301: @*/
1302: PetscErrorCode TSThetaGetTheta(TS ts, PetscReal *theta)
1303: {
1304:   PetscFunctionBegin;
1306:   PetscAssertPointer(theta, 2);
1307:   PetscUseMethod(ts, "TSThetaGetTheta_C", (TS, PetscReal *), (ts, theta));
1308:   PetscFunctionReturn(PETSC_SUCCESS);
1309: }

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

1314:   Not Collective

1316:   Input Parameters:
1317: + ts    - timestepping context
1318: - theta - stage abscissa

1320:   Options Database Key:
1321: . -ts_theta_theta <theta> - set theta

1323:   Level: intermediate

1325: .seealso: [](ch_ts), `TSThetaGetTheta()`, `TSTHETA`, `TSCN`
1326: @*/
1327: PetscErrorCode TSThetaSetTheta(TS ts, PetscReal theta)
1328: {
1329:   PetscFunctionBegin;
1331:   PetscTryMethod(ts, "TSThetaSetTheta_C", (TS, PetscReal), (ts, theta));
1332:   PetscFunctionReturn(PETSC_SUCCESS);
1333: }

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

1338:   Not Collective

1340:   Input Parameter:
1341: . ts - timestepping context

1343:   Output Parameter:
1344: . endpoint - `PETSC_TRUE` when using the endpoint variant

1346:   Level: advanced

1348: .seealso: [](ch_ts), `TSThetaSetEndpoint()`, `TSTHETA`, `TSCN`
1349: @*/
1350: PetscErrorCode TSThetaGetEndpoint(TS ts, PetscBool *endpoint)
1351: {
1352:   PetscFunctionBegin;
1354:   PetscAssertPointer(endpoint, 2);
1355:   PetscUseMethod(ts, "TSThetaGetEndpoint_C", (TS, PetscBool *), (ts, endpoint));
1356:   PetscFunctionReturn(PETSC_SUCCESS);
1357: }

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

1362:   Not Collective

1364:   Input Parameters:
1365: + ts  - timestepping context
1366: - flg - `PETSC_TRUE` to use the endpoint variant

1368:   Options Database Key:
1369: . -ts_theta_endpoint <flg> - use the endpoint variant

1371:   Level: intermediate

1373: .seealso: [](ch_ts), `TSTHETA`, `TSCN`
1374: @*/
1375: PetscErrorCode TSThetaSetEndpoint(TS ts, PetscBool flg)
1376: {
1377:   PetscFunctionBegin;
1379:   PetscTryMethod(ts, "TSThetaSetEndpoint_C", (TS, PetscBool), (ts, flg));
1380:   PetscFunctionReturn(PETSC_SUCCESS);
1381: }

1383: /*
1384:  * TSBEULER and TSCN are straightforward specializations of TSTHETA.
1385:  * The creation functions for these specializations are below.
1386:  */

1388: static PetscErrorCode TSSetUp_BEuler(TS ts)
1389: {
1390:   TS_Theta *th = (TS_Theta *)ts->data;

1392:   PetscFunctionBegin;
1393:   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");
1394:   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");
1395:   PetscCall(TSSetUp_Theta(ts));
1396:   PetscFunctionReturn(PETSC_SUCCESS);
1397: }

1399: static PetscErrorCode TSView_BEuler(TS ts, PetscViewer viewer)
1400: {
1401:   PetscFunctionBegin;
1402:   PetscFunctionReturn(PETSC_SUCCESS);
1403: }

1405: /*MC
1406:       TSBEULER - ODE solver using the implicit backward Euler method

1408:   Level: beginner

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

1413: .seealso: [](ch_ts), `TSCreate()`, `TS`, `TSSetType()`, `TSEULER`, `TSCN`, `TSTHETA`
1414: M*/
1415: PETSC_EXTERN PetscErrorCode TSCreate_BEuler(TS ts)
1416: {
1417:   PetscFunctionBegin;
1418:   PetscCall(TSCreate_Theta(ts));
1419:   PetscCall(TSThetaSetTheta(ts, 1.0));
1420:   PetscCall(TSThetaSetEndpoint(ts, PETSC_FALSE));
1421:   ts->ops->setup = TSSetUp_BEuler;
1422:   ts->ops->view  = TSView_BEuler;
1423:   PetscFunctionReturn(PETSC_SUCCESS);
1424: }

1426: static PetscErrorCode TSSetUp_CN(TS ts)
1427: {
1428:   TS_Theta *th = (TS_Theta *)ts->data;

1430:   PetscFunctionBegin;
1431:   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");
1432:   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");
1433:   PetscCall(TSSetUp_Theta(ts));
1434:   PetscFunctionReturn(PETSC_SUCCESS);
1435: }

1437: static PetscErrorCode TSView_CN(TS ts, PetscViewer viewer)
1438: {
1439:   PetscFunctionBegin;
1440:   PetscFunctionReturn(PETSC_SUCCESS);
1441: }

1443: /*MC
1444:       TSCN - ODE solver using the implicit Crank-Nicolson method.

1446:   Level: beginner

1448:   Notes:
1449:   `TSCN` is equivalent to `TSTHETA` with Theta=0.5 and the "endpoint" option set. I.e.
1450: .vb
1451:   -ts_type theta
1452:   -ts_theta_theta 0.5
1453:   -ts_theta_endpoint
1454: .ve

1456: .seealso: [](ch_ts), `TSCreate()`, `TS`, `TSSetType()`, `TSBEULER`, `TSTHETA`, `TSType`,
1457: M*/
1458: PETSC_EXTERN PetscErrorCode TSCreate_CN(TS ts)
1459: {
1460:   PetscFunctionBegin;
1461:   PetscCall(TSCreate_Theta(ts));
1462:   PetscCall(TSThetaSetTheta(ts, 0.5));
1463:   PetscCall(TSThetaSetEndpoint(ts, PETSC_TRUE));
1464:   ts->ops->setup = TSSetUp_CN;
1465:   ts->ops->view  = TSView_CN;
1466:   PetscFunctionReturn(PETSC_SUCCESS);
1467: }