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 TSStep_Theta(TS ts)
183: {
184:   TS_Theta *th         = (TS_Theta *)ts->data;
185:   PetscInt  rejections = 0;
186:   PetscBool stageok, accept = PETSC_TRUE;
187:   PetscReal next_time_step = ts->time_step;

189:   PetscFunctionBegin;
190:   if (!ts->steprollback) {
191:     if (th->vec_sol_prev) PetscCall(VecCopy(th->X0, th->vec_sol_prev));
192:     PetscCall(VecCopy(ts->vec_sol, th->X0));
193:   }

195:   th->status = TS_STEP_INCOMPLETE;
196:   while (!ts->reason && th->status != TS_STEP_COMPLETE) {
197:     th->shift      = 1 / (th->Theta * ts->time_step);
198:     th->stage_time = ts->ptime + (th->endpoint ? (PetscReal)1 : th->Theta) * ts->time_step;
199:     PetscCall(VecCopy(th->X0, th->X));
200:     if (th->extrapolate && !ts->steprestart) PetscCall(VecAXPY(th->X, 1 / th->shift, th->Xdot));
201:     if (th->endpoint) { /* This formulation assumes linear time-independent mass matrix */
202:       if (!th->affine) PetscCall(VecDuplicate(ts->vec_sol, &th->affine));
203:       PetscCall(VecZeroEntries(th->Xdot));
204:       PetscCall(TSComputeIFunction(ts, ts->ptime, th->X0, th->Xdot, th->affine, PETSC_FALSE));
205:       PetscCall(VecScale(th->affine, (th->Theta - 1) / th->Theta));
206:     } else if (th->affine) { /* Just in case th->endpoint is changed between calls to TSStep_Theta() */
207:       PetscCall(VecZeroEntries(th->affine));
208:     }
209:     PetscCall(TSPreStage(ts, th->stage_time));
210:     PetscCall(TSTheta_SNESSolve(ts, th->affine, th->X));
211:     PetscCall(TSPostStage(ts, th->stage_time, 0, &th->X));
212:     PetscCall(TSAdaptCheckStage(ts->adapt, ts, th->stage_time, th->X, &stageok));
213:     if (!stageok) goto reject_step;

215:     th->status = TS_STEP_PENDING;
216:     if (th->endpoint) {
217:       PetscCall(VecCopy(th->X, ts->vec_sol));
218:     } else {
219:       PetscCall(VecAXPBYPCZ(th->Xdot, -th->shift, th->shift, 0, th->X0, th->X));
220:       PetscCall(VecAXPY(ts->vec_sol, ts->time_step, th->Xdot));
221:     }
222:     PetscCall(TSAdaptChoose(ts->adapt, ts, ts->time_step, NULL, &next_time_step, &accept));
223:     th->status = accept ? TS_STEP_COMPLETE : TS_STEP_INCOMPLETE;
224:     if (!accept) {
225:       PetscCall(VecCopy(th->X0, ts->vec_sol));
226:       ts->time_step = next_time_step;
227:       goto reject_step;
228:     }

230:     if (ts->forward_solve || ts->costintegralfwd) { /* Save the info for the later use in cost integral evaluation */
231:       th->ptime0     = ts->ptime;
232:       th->time_step0 = ts->time_step;
233:     }
234:     ts->ptime += ts->time_step;
235:     ts->time_step = next_time_step;
236:     break;

238:   reject_step:
239:     ts->reject++;
240:     accept = PETSC_FALSE;
241:     if (!ts->reason && ++rejections > ts->max_reject && ts->max_reject >= 0) {
242:       ts->reason = TS_DIVERGED_STEP_REJECTED;
243:       PetscCall(PetscInfo(ts, "Step=%" PetscInt_FMT ", step rejections %" PetscInt_FMT " greater than current TS allowed, stopping solve\n", ts->steps, rejections));
244:     }
245:   }
246:   PetscFunctionReturn(PETSC_SUCCESS);
247: }

249: static PetscErrorCode TSAdjointStepBEuler_Private(TS ts)
250: {
251:   TS_Theta      *th           = (TS_Theta *)ts->data;
252:   TS             quadts       = ts->quadraturets;
253:   Vec           *VecsDeltaLam = th->VecsDeltaLam, *VecsDeltaMu = th->VecsDeltaMu, *VecsSensiTemp = th->VecsSensiTemp;
254:   Vec           *VecsDeltaLam2 = th->VecsDeltaLam2, *VecsDeltaMu2 = th->VecsDeltaMu2, *VecsSensi2Temp = th->VecsSensi2Temp;
255:   PetscInt       nadj;
256:   Mat            J, Jpre, quadJ = NULL, quadJp = NULL;
257:   KSP            ksp;
258:   PetscScalar   *xarr;
259:   TSEquationType eqtype;
260:   PetscBool      isexplicitode = PETSC_FALSE;
261:   PetscReal      adjoint_time_step;

263:   PetscFunctionBegin;
264:   PetscCall(TSGetEquationType(ts, &eqtype));
265:   if (eqtype == TS_EQ_ODE_EXPLICIT) {
266:     isexplicitode = PETSC_TRUE;
267:     VecsDeltaLam  = ts->vecs_sensi;
268:     VecsDeltaLam2 = ts->vecs_sensi2;
269:   }
270:   th->status = TS_STEP_INCOMPLETE;
271:   PetscCall(SNESGetKSP(ts->snes, &ksp));
272:   PetscCall(TSGetIJacobian(ts, &J, &Jpre, NULL, NULL));
273:   if (quadts) {
274:     PetscCall(TSGetRHSJacobian(quadts, &quadJ, NULL, NULL, NULL));
275:     PetscCall(TSGetRHSJacobianP(quadts, &quadJp, NULL, NULL));
276:   }

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

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

284:   for (nadj = 0; nadj < ts->numcost; nadj++) {
285:     PetscCall(VecCopy(ts->vecs_sensi[nadj], VecsSensiTemp[nadj]));
286:     PetscCall(VecScale(VecsSensiTemp[nadj], 1. / adjoint_time_step)); /* lambda_{n+1}/h */
287:     if (quadJ) {
288:       PetscCall(MatDenseGetColumn(quadJ, nadj, &xarr));
289:       PetscCall(VecPlaceArray(ts->vec_drdu_col, xarr));
290:       PetscCall(VecAXPY(VecsSensiTemp[nadj], 1., ts->vec_drdu_col));
291:       PetscCall(VecResetArray(ts->vec_drdu_col));
292:       PetscCall(MatDenseRestoreColumn(quadJ, &xarr));
293:     }
294:   }

296:   /* Build LHS for first-order adjoint */
297:   th->shift = 1. / adjoint_time_step;
298:   PetscCall(TSComputeSNESJacobian(ts, ts->vec_sol, J, Jpre));
299:   PetscCall(KSPSetOperators(ksp, J, Jpre));

301:   /* Solve stage equation LHS*lambda_s = RHS for first-order adjoint */
302:   for (nadj = 0; nadj < ts->numcost; nadj++) {
303:     KSPConvergedReason kspreason;
304:     PetscCall(KSPSolveTranspose(ksp, VecsSensiTemp[nadj], VecsDeltaLam[nadj]));
305:     PetscCall(KSPGetConvergedReason(ksp, &kspreason));
306:     if (kspreason < 0) {
307:       ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
308:       PetscCall(PetscInfo(ts, "Step=%" PetscInt_FMT ", %" PetscInt_FMT "th cost function, transposed linear solve fails, stopping 1st-order adjoint solve\n", ts->steps, nadj));
309:     }
310:   }

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

338:   /* Update sensitivities, and evaluate integrals if there is any */
339:   if (!isexplicitode) {
340:     th->shift = 0.0;
341:     PetscCall(TSComputeSNESJacobian(ts, ts->vec_sol, J, Jpre));
342:     PetscCall(KSPSetOperators(ksp, J, Jpre));
343:     for (nadj = 0; nadj < ts->numcost; nadj++) {
344:       /* Add f_U \lambda_s to the original RHS */
345:       PetscCall(VecScale(VecsSensiTemp[nadj], -1.));
346:       PetscCall(MatMultTransposeAdd(J, VecsDeltaLam[nadj], VecsSensiTemp[nadj], VecsSensiTemp[nadj]));
347:       PetscCall(VecScale(VecsSensiTemp[nadj], -adjoint_time_step));
348:       PetscCall(VecCopy(VecsSensiTemp[nadj], ts->vecs_sensi[nadj]));
349:       if (ts->vecs_sensi2) {
350:         PetscCall(MatMultTransposeAdd(J, VecsDeltaLam2[nadj], VecsSensi2Temp[nadj], VecsSensi2Temp[nadj]));
351:         PetscCall(VecScale(VecsSensi2Temp[nadj], -adjoint_time_step));
352:         PetscCall(VecCopy(VecsSensi2Temp[nadj], ts->vecs_sensi2[nadj]));
353:       }
354:     }
355:   }
356:   if (ts->vecs_sensip) {
357:     PetscCall(VecAXPBYPCZ(th->Xdot, -1. / adjoint_time_step, 1.0 / adjoint_time_step, 0, th->X0, ts->vec_sol));
358:     PetscCall(TSComputeIJacobianP(ts, th->stage_time, ts->vec_sol, th->Xdot, 1. / adjoint_time_step, ts->Jacp, PETSC_FALSE)); /* get -f_p */
359:     if (quadts) PetscCall(TSComputeRHSJacobianP(quadts, th->stage_time, ts->vec_sol, quadJp));
360:     if (ts->vecs_sensi2p) {
361:       /* lambda_s^T F_PU w_1 */
362:       PetscCall(TSComputeIHessianProductFunctionPU(ts, th->stage_time, ts->vec_sol, VecsDeltaLam, ts->vec_sensip_col, ts->vecs_fpu));
363:       /* lambda_s^T F_PP w_2 */
364:       PetscCall(TSComputeIHessianProductFunctionPP(ts, th->stage_time, ts->vec_sol, VecsDeltaLam, ts->vec_dir, ts->vecs_fpp));
365:     }

367:     for (nadj = 0; nadj < ts->numcost; nadj++) {
368:       PetscCall(MatMultTranspose(ts->Jacp, VecsDeltaLam[nadj], VecsDeltaMu[nadj]));
369:       PetscCall(VecAXPY(ts->vecs_sensip[nadj], -adjoint_time_step, VecsDeltaMu[nadj]));
370:       if (quadJp) {
371:         PetscCall(MatDenseGetColumn(quadJp, nadj, &xarr));
372:         PetscCall(VecPlaceArray(ts->vec_drdp_col, xarr));
373:         PetscCall(VecAXPY(ts->vecs_sensip[nadj], adjoint_time_step, ts->vec_drdp_col));
374:         PetscCall(VecResetArray(ts->vec_drdp_col));
375:         PetscCall(MatDenseRestoreColumn(quadJp, &xarr));
376:       }
377:       if (ts->vecs_sensi2p) {
378:         PetscCall(MatMultTranspose(ts->Jacp, VecsDeltaLam2[nadj], VecsDeltaMu2[nadj]));
379:         PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step, VecsDeltaMu2[nadj]));
380:         if (ts->vecs_fpu) PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step, ts->vecs_fpu[nadj]));
381:         if (ts->vecs_fpp) PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step, ts->vecs_fpp[nadj]));
382:       }
383:     }
384:   }

386:   if (ts->vecs_sensi2) {
387:     PetscCall(VecResetArray(ts->vec_sensip_col));
388:     PetscCall(MatDenseRestoreColumn(ts->mat_sensip, &xarr));
389:   }
390:   th->status = TS_STEP_COMPLETE;
391:   PetscFunctionReturn(PETSC_SUCCESS);
392: }

394: static PetscErrorCode TSAdjointStep_Theta(TS ts)
395: {
396:   TS_Theta    *th           = (TS_Theta *)ts->data;
397:   TS           quadts       = ts->quadraturets;
398:   Vec         *VecsDeltaLam = th->VecsDeltaLam, *VecsDeltaMu = th->VecsDeltaMu, *VecsSensiTemp = th->VecsSensiTemp;
399:   Vec         *VecsDeltaLam2 = th->VecsDeltaLam2, *VecsDeltaMu2 = th->VecsDeltaMu2, *VecsSensi2Temp = th->VecsSensi2Temp;
400:   PetscInt     nadj;
401:   Mat          J, Jpre, quadJ = NULL, quadJp = NULL;
402:   KSP          ksp;
403:   PetscScalar *xarr;
404:   PetscReal    adjoint_time_step;
405:   PetscReal    adjoint_ptime; /* end time of the adjoint time step (ts->ptime is the start time, usually ts->ptime is larger than adjoint_ptime) */

407:   PetscFunctionBegin;
408:   if (th->Theta == 1.) {
409:     PetscCall(TSAdjointStepBEuler_Private(ts));
410:     PetscFunctionReturn(PETSC_SUCCESS);
411:   }
412:   th->status = TS_STEP_INCOMPLETE;
413:   PetscCall(SNESGetKSP(ts->snes, &ksp));
414:   PetscCall(TSGetIJacobian(ts, &J, &Jpre, NULL, NULL));
415:   if (quadts) {
416:     PetscCall(TSGetRHSJacobian(quadts, &quadJ, NULL, NULL, NULL));
417:     PetscCall(TSGetRHSJacobianP(quadts, &quadJp, NULL, NULL));
418:   }
419:   /* If endpoint=1, th->ptime and th->X0 will be used; if endpoint=0, th->stage_time and th->X will be used. */
420:   th->stage_time    = th->endpoint ? ts->ptime : (ts->ptime + (1. - th->Theta) * ts->time_step);
421:   adjoint_ptime     = ts->ptime + ts->time_step;
422:   adjoint_time_step = -ts->time_step; /* always positive since time_step is negative */

424:   if (!th->endpoint) {
425:     /* recover th->X0 using vec_sol and the stage value th->X */
426:     PetscCall(VecAXPBYPCZ(th->X0, 1.0 / (1.0 - th->Theta), th->Theta / (th->Theta - 1.0), 0, th->X, ts->vec_sol));
427:   }

429:   /* Build RHS for first-order adjoint */
430:   /* Cost function has an integral term */
431:   if (quadts) {
432:     if (th->endpoint) {
433:       PetscCall(TSComputeRHSJacobian(quadts, th->stage_time, ts->vec_sol, quadJ, NULL));
434:     } else {
435:       PetscCall(TSComputeRHSJacobian(quadts, th->stage_time, th->X, quadJ, NULL));
436:     }
437:   }

439:   for (nadj = 0; nadj < ts->numcost; nadj++) {
440:     PetscCall(VecCopy(ts->vecs_sensi[nadj], VecsSensiTemp[nadj]));
441:     PetscCall(VecScale(VecsSensiTemp[nadj], 1. / (th->Theta * adjoint_time_step)));
442:     if (quadJ) {
443:       PetscCall(MatDenseGetColumn(quadJ, nadj, &xarr));
444:       PetscCall(VecPlaceArray(ts->vec_drdu_col, xarr));
445:       PetscCall(VecAXPY(VecsSensiTemp[nadj], 1., ts->vec_drdu_col));
446:       PetscCall(VecResetArray(ts->vec_drdu_col));
447:       PetscCall(MatDenseRestoreColumn(quadJ, &xarr));
448:     }
449:   }

451:   /* Build LHS for first-order adjoint */
452:   th->shift = 1. / (th->Theta * adjoint_time_step);
453:   if (th->endpoint) {
454:     PetscCall(TSComputeSNESJacobian(ts, ts->vec_sol, J, Jpre));
455:   } else {
456:     PetscCall(TSComputeSNESJacobian(ts, th->X, J, Jpre));
457:   }
458:   PetscCall(KSPSetOperators(ksp, J, Jpre));

460:   /* Solve stage equation LHS*lambda_s = RHS for first-order adjoint */
461:   for (nadj = 0; nadj < ts->numcost; nadj++) {
462:     KSPConvergedReason kspreason;
463:     PetscCall(KSPSolveTranspose(ksp, VecsSensiTemp[nadj], VecsDeltaLam[nadj]));
464:     PetscCall(KSPGetConvergedReason(ksp, &kspreason));
465:     if (kspreason < 0) {
466:       ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
467:       PetscCall(PetscInfo(ts, "Step=%" PetscInt_FMT ", %" PetscInt_FMT "th cost function, transposed linear solve fails, stopping 1st-order adjoint solve\n", ts->steps, nadj));
468:     }
469:   }

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

501:   /* Update sensitivities, and evaluate integrals if there is any */
502:   if (th->endpoint) { /* two-stage Theta methods with th->Theta!=1, th->Theta==1 leads to BEuler */
503:     th->shift      = 1. / ((th->Theta - 1.) * adjoint_time_step);
504:     th->stage_time = adjoint_ptime;
505:     PetscCall(TSComputeSNESJacobian(ts, th->X0, J, Jpre));
506:     PetscCall(KSPSetOperators(ksp, J, Jpre));
507:     /* R_U at t_n */
508:     if (quadts) PetscCall(TSComputeRHSJacobian(quadts, adjoint_ptime, th->X0, quadJ, NULL));
509:     for (nadj = 0; nadj < ts->numcost; nadj++) {
510:       PetscCall(MatMultTranspose(J, VecsDeltaLam[nadj], ts->vecs_sensi[nadj]));
511:       if (quadJ) {
512:         PetscCall(MatDenseGetColumn(quadJ, nadj, &xarr));
513:         PetscCall(VecPlaceArray(ts->vec_drdu_col, xarr));
514:         PetscCall(VecAXPY(ts->vecs_sensi[nadj], -1., ts->vec_drdu_col));
515:         PetscCall(VecResetArray(ts->vec_drdu_col));
516:         PetscCall(MatDenseRestoreColumn(quadJ, &xarr));
517:       }
518:       PetscCall(VecScale(ts->vecs_sensi[nadj], 1. / th->shift));
519:     }

521:     /* Second-order adjoint */
522:     if (ts->vecs_sensi2) { /* U_n */
523:       /* Get w1 at t_n from TLM matrix */
524:       PetscCall(MatDenseGetColumn(th->MatFwdSensip0, 0, &xarr));
525:       PetscCall(VecPlaceArray(ts->vec_sensip_col, xarr));
526:       /* lambda_s^T F_UU w_1 */
527:       PetscCall(TSComputeIHessianProductFunctionUU(ts, adjoint_ptime, th->X0, VecsDeltaLam, ts->vec_sensip_col, ts->vecs_fuu));
528:       PetscCall(VecResetArray(ts->vec_sensip_col));
529:       PetscCall(MatDenseRestoreColumn(th->MatFwdSensip0, &xarr));
530:       /* lambda_s^T F_UU w_2 */
531:       PetscCall(TSComputeIHessianProductFunctionUP(ts, adjoint_ptime, th->X0, VecsDeltaLam, ts->vec_dir, ts->vecs_fup));
532:       for (nadj = 0; nadj < ts->numcost; nadj++) {
533:         /* 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  */
534:         PetscCall(MatMultTranspose(J, VecsDeltaLam2[nadj], ts->vecs_sensi2[nadj]));
535:         PetscCall(VecAXPY(ts->vecs_sensi2[nadj], 1., ts->vecs_fuu[nadj]));
536:         if (ts->vecs_fup) PetscCall(VecAXPY(ts->vecs_sensi2[nadj], 1., ts->vecs_fup[nadj]));
537:         PetscCall(VecScale(ts->vecs_sensi2[nadj], 1. / th->shift));
538:       }
539:     }

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

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

569:         /* lambda_s^T F_PP w_2 */
570:         PetscCall(TSComputeIHessianProductFunctionPP(ts, th->stage_time, ts->vec_sol, VecsDeltaLam, ts->vec_dir, ts->vecs_fpp));
571:         for (nadj = 0; nadj < ts->numcost; nadj++) {
572:           /* 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)  */
573:           PetscCall(MatMultTranspose(ts->Jacp, VecsDeltaLam2[nadj], VecsDeltaMu2[nadj]));
574:           PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step * th->Theta, VecsDeltaMu2[nadj]));
575:           if (ts->vecs_fpu) PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step * th->Theta, ts->vecs_fpu[nadj]));
576:           if (ts->vecs_fpp) PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step * th->Theta, ts->vecs_fpp[nadj]));
577:         }
578:       }

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

649:   th->status = TS_STEP_COMPLETE;
650:   PetscFunctionReturn(PETSC_SUCCESS);
651: }

653: static PetscErrorCode TSInterpolate_Theta(TS ts, PetscReal t, Vec X)
654: {
655:   TS_Theta *th = (TS_Theta *)ts->data;
656:   PetscReal dt = t - ts->ptime;

658:   PetscFunctionBegin;
659:   PetscCall(VecCopy(ts->vec_sol, th->X));
660:   if (th->endpoint) dt *= th->Theta;
661:   PetscCall(VecWAXPY(X, dt, th->Xdot, th->X));
662:   PetscFunctionReturn(PETSC_SUCCESS);
663: }

665: static PetscErrorCode TSEvaluateWLTE_Theta(TS ts, NormType wnormtype, PetscInt *order, PetscReal *wlte)
666: {
667:   TS_Theta *th = (TS_Theta *)ts->data;
668:   Vec       X  = ts->vec_sol;      /* X = solution */
669:   Vec       Y  = th->vec_lte_work; /* Y = X + LTE  */
670:   PetscReal wltea, wlter;

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

702: static PetscErrorCode TSRollBack_Theta(TS ts)
703: {
704:   TS_Theta *th     = (TS_Theta *)ts->data;
705:   TS        quadts = ts->quadraturets;

707:   PetscFunctionBegin;
708:   PetscCall(VecCopy(th->X0, ts->vec_sol));
709:   if (quadts && ts->costintegralfwd) PetscCall(VecCopy(th->VecCostIntegral0, quadts->vec_sol));
710:   th->status = TS_STEP_INCOMPLETE;
711:   if (ts->mat_sensip) PetscCall(MatCopy(th->MatFwdSensip0, ts->mat_sensip, SAME_NONZERO_PATTERN));
712:   if (quadts && quadts->mat_sensip) PetscCall(MatCopy(th->MatIntegralSensip0, quadts->mat_sensip, SAME_NONZERO_PATTERN));
713:   PetscFunctionReturn(PETSC_SUCCESS);
714: }

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

728:   PetscFunctionBegin;
729:   previous_shift = th->shift;
730:   PetscCall(MatCopy(ts->mat_sensip, th->MatFwdSensip0, SAME_NONZERO_PATTERN));

732:   if (quadts && quadts->mat_sensip) PetscCall(MatCopy(quadts->mat_sensip, th->MatIntegralSensip0, SAME_NONZERO_PATTERN));
733:   PetscCall(SNESGetKSP(ts->snes, &ksp));
734:   PetscCall(TSGetIJacobian(ts, &J, &Jpre, NULL, NULL));
735:   if (quadts) {
736:     PetscCall(TSGetRHSJacobian(quadts, &quadJ, NULL, NULL, NULL));
737:     PetscCall(TSGetRHSJacobianP(quadts, &quadJp, NULL, NULL));
738:   }

740:   /* Build RHS */
741:   if (th->endpoint) { /* 2-stage method*/
742:     th->shift = 1. / ((th->Theta - 1.) * th->time_step0);
743:     PetscCall(TSComputeIJacobian(ts, th->ptime0, th->X0, th->Xdot, th->shift, J, Jpre, PETSC_FALSE));
744:     PetscCall(MatMatMult(J, ts->mat_sensip, MAT_REUSE_MATRIX, PETSC_DEFAULT, &MatDeltaFwdSensip));
745:     PetscCall(MatScale(MatDeltaFwdSensip, (th->Theta - 1.) / th->Theta));

747:     /* Add the f_p forcing terms */
748:     if (ts->Jacp) {
749:       PetscCall(VecZeroEntries(th->Xdot));
750:       PetscCall(TSComputeIJacobianP(ts, th->ptime0, th->X0, th->Xdot, th->shift, ts->Jacp, PETSC_FALSE));
751:       PetscCall(MatAXPY(MatDeltaFwdSensip, (th->Theta - 1.) / th->Theta, ts->Jacp, SUBSET_NONZERO_PATTERN));
752:       th->shift = previous_shift;
753:       PetscCall(VecAXPBYPCZ(th->Xdot, -th->shift, th->shift, 0, th->X0, ts->vec_sol));
754:       PetscCall(TSComputeIJacobianP(ts, th->stage_time, ts->vec_sol, th->Xdot, th->shift, ts->Jacp, PETSC_FALSE));
755:       PetscCall(MatAXPY(MatDeltaFwdSensip, -1., ts->Jacp, SUBSET_NONZERO_PATTERN));
756:     }
757:   } else { /* 1-stage method */
758:     th->shift = 0.0;
759:     PetscCall(TSComputeIJacobian(ts, th->stage_time, th->X, th->Xdot, th->shift, J, Jpre, PETSC_FALSE));
760:     PetscCall(MatMatMult(J, ts->mat_sensip, MAT_REUSE_MATRIX, PETSC_DEFAULT, &MatDeltaFwdSensip));
761:     PetscCall(MatScale(MatDeltaFwdSensip, -1.));

763:     /* Add the f_p forcing terms */
764:     if (ts->Jacp) {
765:       th->shift = previous_shift;
766:       PetscCall(VecAXPBYPCZ(th->Xdot, -th->shift, th->shift, 0, th->X0, th->X));
767:       PetscCall(TSComputeIJacobianP(ts, th->stage_time, th->X, th->Xdot, th->shift, ts->Jacp, PETSC_FALSE));
768:       PetscCall(MatAXPY(MatDeltaFwdSensip, -1., ts->Jacp, SUBSET_NONZERO_PATTERN));
769:     }
770:   }

772:   /* Build LHS */
773:   th->shift = previous_shift; /* recover the previous shift used in TSStep_Theta() */
774:   if (th->endpoint) {
775:     PetscCall(TSComputeIJacobian(ts, th->stage_time, ts->vec_sol, th->Xdot, th->shift, J, Jpre, PETSC_FALSE));
776:   } else {
777:     PetscCall(TSComputeIJacobian(ts, th->stage_time, th->X, th->Xdot, th->shift, J, Jpre, PETSC_FALSE));
778:   }
779:   PetscCall(KSPSetOperators(ksp, J, Jpre));

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

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

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

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

849:   PetscFunctionBegin;
850:   if (ns) {
851:     if (!th->endpoint && th->Theta != 1.0) *ns = 1; /* midpoint form */
852:     else *ns = 2;                                   /* endpoint form */
853:   }
854:   if (stagesensip) {
855:     if (!th->endpoint && th->Theta != 1.0) {
856:       th->MatFwdStages[0] = th->MatDeltaFwdSensip;
857:     } else {
858:       th->MatFwdStages[0] = th->MatFwdSensip0;
859:       th->MatFwdStages[1] = ts->mat_sensip; /* stiffly accurate */
860:     }
861:     *stagesensip = th->MatFwdStages;
862:   }
863:   PetscFunctionReturn(PETSC_SUCCESS);
864: }

866: /*------------------------------------------------------------*/
867: static PetscErrorCode TSReset_Theta(TS ts)
868: {
869:   TS_Theta *th = (TS_Theta *)ts->data;

871:   PetscFunctionBegin;
872:   PetscCall(VecDestroy(&th->X));
873:   PetscCall(VecDestroy(&th->Xdot));
874:   PetscCall(VecDestroy(&th->X0));
875:   PetscCall(VecDestroy(&th->affine));

877:   PetscCall(VecDestroy(&th->vec_sol_prev));
878:   PetscCall(VecDestroy(&th->vec_lte_work));

880:   PetscCall(VecDestroy(&th->VecCostIntegral0));
881:   PetscFunctionReturn(PETSC_SUCCESS);
882: }

884: static PetscErrorCode TSAdjointReset_Theta(TS ts)
885: {
886:   TS_Theta *th = (TS_Theta *)ts->data;

888:   PetscFunctionBegin;
889:   PetscCall(VecDestroyVecs(ts->numcost, &th->VecsDeltaLam));
890:   PetscCall(VecDestroyVecs(ts->numcost, &th->VecsDeltaMu));
891:   PetscCall(VecDestroyVecs(ts->numcost, &th->VecsDeltaLam2));
892:   PetscCall(VecDestroyVecs(ts->numcost, &th->VecsDeltaMu2));
893:   PetscCall(VecDestroyVecs(ts->numcost, &th->VecsSensiTemp));
894:   PetscCall(VecDestroyVecs(ts->numcost, &th->VecsSensi2Temp));
895:   PetscFunctionReturn(PETSC_SUCCESS);
896: }

898: static PetscErrorCode TSDestroy_Theta(TS ts)
899: {
900:   PetscFunctionBegin;
901:   PetscCall(TSReset_Theta(ts));
902:   if (ts->dm) {
903:     PetscCall(DMCoarsenHookRemove(ts->dm, DMCoarsenHook_TSTheta, DMRestrictHook_TSTheta, ts));
904:     PetscCall(DMSubDomainHookRemove(ts->dm, DMSubDomainHook_TSTheta, DMSubDomainRestrictHook_TSTheta, ts));
905:   }
906:   PetscCall(PetscFree(ts->data));
907:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaGetTheta_C", NULL));
908:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaSetTheta_C", NULL));
909:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaGetEndpoint_C", NULL));
910:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaSetEndpoint_C", NULL));
911:   PetscFunctionReturn(PETSC_SUCCESS);
912: }

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

918:   Note that U here is the stage argument. This means that U = U_{n+1} only if endpoint = true,
919:   otherwise U = theta U_{n+1} + (1 - theta) U0, which for the case of implicit midpoint is
920:   U = (U_{n+1} + U0)/2
921: */
922: static PetscErrorCode SNESTSFormFunction_Theta(SNES snes, Vec x, Vec y, TS ts)
923: {
924:   TS_Theta *th = (TS_Theta *)ts->data;
925:   Vec       X0, Xdot;
926:   DM        dm, dmsave;
927:   PetscReal shift = th->shift;

929:   PetscFunctionBegin;
930:   PetscCall(SNESGetDM(snes, &dm));
931:   /* When using the endpoint variant, this is actually 1/Theta * Xdot */
932:   PetscCall(TSThetaGetX0AndXdot(ts, dm, &X0, &Xdot));
933:   if (x != X0) {
934:     PetscCall(VecAXPBYPCZ(Xdot, -shift, shift, 0, X0, x));
935:   } else {
936:     PetscCall(VecZeroEntries(Xdot));
937:   }
938:   /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */
939:   dmsave = ts->dm;
940:   ts->dm = dm;
941:   PetscCall(TSComputeIFunction(ts, th->stage_time, x, Xdot, y, PETSC_FALSE));
942:   ts->dm = dmsave;
943:   PetscCall(TSThetaRestoreX0AndXdot(ts, dm, &X0, &Xdot));
944:   PetscFunctionReturn(PETSC_SUCCESS);
945: }

947: static PetscErrorCode SNESTSFormJacobian_Theta(SNES snes, Vec x, Mat A, Mat B, TS ts)
948: {
949:   TS_Theta *th = (TS_Theta *)ts->data;
950:   Vec       Xdot;
951:   DM        dm, dmsave;
952:   PetscReal shift = th->shift;

954:   PetscFunctionBegin;
955:   PetscCall(SNESGetDM(snes, &dm));
956:   /* Xdot has already been computed in SNESTSFormFunction_Theta (SNES guarantees this) */
957:   PetscCall(TSThetaGetX0AndXdot(ts, dm, NULL, &Xdot));

959:   dmsave = ts->dm;
960:   ts->dm = dm;
961:   PetscCall(TSComputeIJacobian(ts, th->stage_time, x, Xdot, shift, A, B, PETSC_FALSE));
962:   ts->dm = dmsave;
963:   PetscCall(TSThetaRestoreX0AndXdot(ts, dm, NULL, &Xdot));
964:   PetscFunctionReturn(PETSC_SUCCESS);
965: }

967: static PetscErrorCode TSForwardSetUp_Theta(TS ts)
968: {
969:   TS_Theta *th     = (TS_Theta *)ts->data;
970:   TS        quadts = ts->quadraturets;

972:   PetscFunctionBegin;
973:   /* combine sensitivities to parameters and sensitivities to initial values into one array */
974:   th->num_tlm = ts->num_parameters;
975:   PetscCall(MatDuplicate(ts->mat_sensip, MAT_DO_NOT_COPY_VALUES, &th->MatDeltaFwdSensip));
976:   if (quadts && quadts->mat_sensip) {
977:     PetscCall(MatDuplicate(quadts->mat_sensip, MAT_DO_NOT_COPY_VALUES, &th->MatIntegralSensipTemp));
978:     PetscCall(MatDuplicate(quadts->mat_sensip, MAT_DO_NOT_COPY_VALUES, &th->MatIntegralSensip0));
979:   }
980:   /* backup sensitivity results for roll-backs */
981:   PetscCall(MatDuplicate(ts->mat_sensip, MAT_DO_NOT_COPY_VALUES, &th->MatFwdSensip0));

983:   PetscCall(VecDuplicate(ts->vec_sol, &th->VecDeltaFwdSensipCol));
984:   PetscFunctionReturn(PETSC_SUCCESS);
985: }

987: static PetscErrorCode TSForwardReset_Theta(TS ts)
988: {
989:   TS_Theta *th     = (TS_Theta *)ts->data;
990:   TS        quadts = ts->quadraturets;

992:   PetscFunctionBegin;
993:   if (quadts && quadts->mat_sensip) {
994:     PetscCall(MatDestroy(&th->MatIntegralSensipTemp));
995:     PetscCall(MatDestroy(&th->MatIntegralSensip0));
996:   }
997:   PetscCall(VecDestroy(&th->VecDeltaFwdSensipCol));
998:   PetscCall(MatDestroy(&th->MatDeltaFwdSensip));
999:   PetscCall(MatDestroy(&th->MatFwdSensip0));
1000:   PetscFunctionReturn(PETSC_SUCCESS);
1001: }

1003: static PetscErrorCode TSSetUp_Theta(TS ts)
1004: {
1005:   TS_Theta *th     = (TS_Theta *)ts->data;
1006:   TS        quadts = ts->quadraturets;
1007:   PetscBool match;

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

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

1021:   PetscCall(TSGetDM(ts, &ts->dm));
1022:   PetscCall(DMCoarsenHookAdd(ts->dm, DMCoarsenHook_TSTheta, DMRestrictHook_TSTheta, ts));
1023:   PetscCall(DMSubDomainHookAdd(ts->dm, DMSubDomainHook_TSTheta, DMSubDomainRestrictHook_TSTheta, ts));

1025:   PetscCall(TSGetAdapt(ts, &ts->adapt));
1026:   PetscCall(TSAdaptCandidatesClear(ts->adapt));
1027:   PetscCall(PetscObjectTypeCompare((PetscObject)ts->adapt, TSADAPTNONE, &match));
1028:   if (!match) {
1029:     PetscCall(VecDuplicate(ts->vec_sol, &th->vec_sol_prev));
1030:     PetscCall(VecDuplicate(ts->vec_sol, &th->vec_lte_work));
1031:   }
1032:   PetscCall(TSGetSNES(ts, &ts->snes));

1034:   ts->stifflyaccurate = (!th->endpoint && th->Theta != 1.0) ? PETSC_FALSE : PETSC_TRUE;
1035:   PetscFunctionReturn(PETSC_SUCCESS);
1036: }

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

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

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

1064: static PetscErrorCode TSSetFromOptions_Theta(TS ts, PetscOptionItems *PetscOptionsObject)
1065: {
1066:   TS_Theta *th = (TS_Theta *)ts->data;

1068:   PetscFunctionBegin;
1069:   PetscOptionsHeadBegin(PetscOptionsObject, "Theta ODE solver options");
1070:   {
1071:     PetscCall(PetscOptionsReal("-ts_theta_theta", "Location of stage (0<Theta<=1)", "TSThetaSetTheta", th->Theta, &th->Theta, NULL));
1072:     PetscCall(PetscOptionsBool("-ts_theta_endpoint", "Use the endpoint instead of midpoint form of the Theta method", "TSThetaSetEndpoint", th->endpoint, &th->endpoint, NULL));
1073:     PetscCall(PetscOptionsBool("-ts_theta_initial_guess_extrapolate", "Extrapolate stage initial guess from previous solution (sometimes unstable)", "TSThetaSetExtrapolate", th->extrapolate, &th->extrapolate, NULL));
1074:   }
1075:   PetscOptionsHeadEnd();
1076:   PetscFunctionReturn(PETSC_SUCCESS);
1077: }

1079: static PetscErrorCode TSView_Theta(TS ts, PetscViewer viewer)
1080: {
1081:   TS_Theta *th = (TS_Theta *)ts->data;
1082:   PetscBool iascii;

1084:   PetscFunctionBegin;
1085:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &iascii));
1086:   if (iascii) {
1087:     PetscCall(PetscViewerASCIIPrintf(viewer, "  Theta=%g\n", (double)th->Theta));
1088:     PetscCall(PetscViewerASCIIPrintf(viewer, "  Extrapolation=%s\n", th->extrapolate ? "yes" : "no"));
1089:   }
1090:   PetscFunctionReturn(PETSC_SUCCESS);
1091: }

1093: static PetscErrorCode TSThetaGetTheta_Theta(TS ts, PetscReal *theta)
1094: {
1095:   TS_Theta *th = (TS_Theta *)ts->data;

1097:   PetscFunctionBegin;
1098:   *theta = th->Theta;
1099:   PetscFunctionReturn(PETSC_SUCCESS);
1100: }

1102: static PetscErrorCode TSThetaSetTheta_Theta(TS ts, PetscReal theta)
1103: {
1104:   TS_Theta *th = (TS_Theta *)ts->data;

1106:   PetscFunctionBegin;
1107:   PetscCheck(theta > 0 && theta <= 1, PetscObjectComm((PetscObject)ts), PETSC_ERR_ARG_OUTOFRANGE, "Theta %g not in range (0,1]", (double)theta);
1108:   th->Theta = theta;
1109:   th->order = (th->Theta == 0.5) ? 2 : 1;
1110:   PetscFunctionReturn(PETSC_SUCCESS);
1111: }

1113: static PetscErrorCode TSThetaGetEndpoint_Theta(TS ts, PetscBool *endpoint)
1114: {
1115:   TS_Theta *th = (TS_Theta *)ts->data;

1117:   PetscFunctionBegin;
1118:   *endpoint = th->endpoint;
1119:   PetscFunctionReturn(PETSC_SUCCESS);
1120: }

1122: static PetscErrorCode TSThetaSetEndpoint_Theta(TS ts, PetscBool flg)
1123: {
1124:   TS_Theta *th = (TS_Theta *)ts->data;

1126:   PetscFunctionBegin;
1127:   th->endpoint = flg;
1128:   PetscFunctionReturn(PETSC_SUCCESS);
1129: }

1131: #if defined(PETSC_HAVE_COMPLEX)
1132: static PetscErrorCode TSComputeLinearStability_Theta(TS ts, PetscReal xr, PetscReal xi, PetscReal *yr, PetscReal *yi)
1133: {
1134:   PetscComplex    z   = xr + xi * PETSC_i, f;
1135:   TS_Theta       *th  = (TS_Theta *)ts->data;
1136:   const PetscReal one = 1.0;

1138:   PetscFunctionBegin;
1139:   f   = (one + (one - th->Theta) * z) / (one - th->Theta * z);
1140:   *yr = PetscRealPartComplex(f);
1141:   *yi = PetscImaginaryPartComplex(f);
1142:   PetscFunctionReturn(PETSC_SUCCESS);
1143: }
1144: #endif

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

1150:   PetscFunctionBegin;
1151:   if (ns) {
1152:     if (!th->endpoint && th->Theta != 1.0) *ns = 1; /* midpoint form */
1153:     else *ns = 2;                                   /* endpoint form */
1154:   }
1155:   if (Y) {
1156:     if (!th->endpoint && th->Theta != 1.0) {
1157:       th->Stages[0] = th->X;
1158:     } else {
1159:       th->Stages[0] = th->X0;
1160:       th->Stages[1] = ts->vec_sol; /* stiffly accurate */
1161:     }
1162:     *Y = th->Stages;
1163:   }
1164:   PetscFunctionReturn(PETSC_SUCCESS);
1165: }

1167: /* ------------------------------------------------------------ */
1168: /*MC
1169:       TSTHETA - DAE solver using the implicit Theta method

1171:    Level: beginner

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

1178:    Notes:
1179: .vb
1180:   -ts_type theta -ts_theta_theta 1.0 corresponds to backward Euler (TSBEULER)
1181:   -ts_type theta -ts_theta_theta 0.5 corresponds to the implicit midpoint rule
1182:   -ts_type theta -ts_theta_theta 0.5 -ts_theta_endpoint corresponds to Crank-Nicholson (TSCN)
1183: .ve

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

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

1189: .vb
1190:   Theta | Theta
1191:   -------------
1192:         |  1
1193: .ve

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

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

1199: .vb
1200:   0 | 0         0
1201:   1 | 1-Theta   Theta
1202:   -------------------
1203:     | 1-Theta   Theta
1204: .ve

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

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

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

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

1214: .seealso: [](ch_ts), `TSCreate()`, `TS`, `TSSetType()`, `TSCN`, `TSBEULER`, `TSThetaSetTheta()`, `TSThetaSetEndpoint()`
1215: M*/
1216: PETSC_EXTERN PetscErrorCode TSCreate_Theta(TS ts)
1217: {
1218:   TS_Theta *th;

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

1244:   ts->ops->forwardsetup     = TSForwardSetUp_Theta;
1245:   ts->ops->forwardreset     = TSForwardReset_Theta;
1246:   ts->ops->forwardstep      = TSForwardStep_Theta;
1247:   ts->ops->forwardgetstages = TSForwardGetStages_Theta;

1249:   ts->usessnes = PETSC_TRUE;

1251:   PetscCall(PetscNew(&th));
1252:   ts->data = (void *)th;

1254:   th->VecsDeltaLam   = NULL;
1255:   th->VecsDeltaMu    = NULL;
1256:   th->VecsSensiTemp  = NULL;
1257:   th->VecsSensi2Temp = NULL;

1259:   th->extrapolate = PETSC_FALSE;
1260:   th->Theta       = 0.5;
1261:   th->order       = 2;
1262:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaGetTheta_C", TSThetaGetTheta_Theta));
1263:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaSetTheta_C", TSThetaSetTheta_Theta));
1264:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaGetEndpoint_C", TSThetaGetEndpoint_Theta));
1265:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaSetEndpoint_C", TSThetaSetEndpoint_Theta));
1266:   PetscFunctionReturn(PETSC_SUCCESS);
1267: }

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

1272:   Not Collective

1274:   Input Parameter:
1275: .  ts - timestepping context

1277:   Output Parameter:
1278: .  theta - stage abscissa

1280:   Level: advanced

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

1285: .seealso: [](ch_ts), `TSThetaSetTheta()`, `TSTHETA`
1286: @*/
1287: PetscErrorCode TSThetaGetTheta(TS ts, PetscReal *theta)
1288: {
1289:   PetscFunctionBegin;
1292:   PetscUseMethod(ts, "TSThetaGetTheta_C", (TS, PetscReal *), (ts, theta));
1293:   PetscFunctionReturn(PETSC_SUCCESS);
1294: }

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

1299:   Not Collective

1301:   Input Parameters:
1302: +  ts - timestepping context
1303: -  theta - stage abscissa

1305:   Options Database Key:
1306: .  -ts_theta_theta <theta> - set theta

1308:   Level: intermediate

1310: .seealso: [](ch_ts), `TSThetaGetTheta()`, `TSTHETA`, `TSCN`
1311: @*/
1312: PetscErrorCode TSThetaSetTheta(TS ts, PetscReal theta)
1313: {
1314:   PetscFunctionBegin;
1316:   PetscTryMethod(ts, "TSThetaSetTheta_C", (TS, PetscReal), (ts, theta));
1317:   PetscFunctionReturn(PETSC_SUCCESS);
1318: }

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

1323:   Not Collective

1325:   Input Parameter:
1326: .  ts - timestepping context

1328:   Output Parameter:
1329: .  endpoint - `PETSC_TRUE` when using the endpoint variant

1331:   Level: advanced

1333: .seealso: [](ch_ts), `TSThetaSetEndpoint()`, `TSTHETA`, `TSCN`
1334: @*/
1335: PetscErrorCode TSThetaGetEndpoint(TS ts, PetscBool *endpoint)
1336: {
1337:   PetscFunctionBegin;
1340:   PetscUseMethod(ts, "TSThetaGetEndpoint_C", (TS, PetscBool *), (ts, endpoint));
1341:   PetscFunctionReturn(PETSC_SUCCESS);
1342: }

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

1347:   Not Collective

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

1353:   Options Database Key:
1354: .  -ts_theta_endpoint <flg> - use the endpoint variant

1356:   Level: intermediate

1358: .seealso: [](ch_ts), `TSTHETA`, `TSCN`
1359: @*/
1360: PetscErrorCode TSThetaSetEndpoint(TS ts, PetscBool flg)
1361: {
1362:   PetscFunctionBegin;
1364:   PetscTryMethod(ts, "TSThetaSetEndpoint_C", (TS, PetscBool), (ts, flg));
1365:   PetscFunctionReturn(PETSC_SUCCESS);
1366: }

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

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

1377:   PetscFunctionBegin;
1378:   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");
1379:   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");
1380:   PetscCall(TSSetUp_Theta(ts));
1381:   PetscFunctionReturn(PETSC_SUCCESS);
1382: }

1384: static PetscErrorCode TSView_BEuler(TS ts, PetscViewer viewer)
1385: {
1386:   PetscFunctionBegin;
1387:   PetscFunctionReturn(PETSC_SUCCESS);
1388: }

1390: /*MC
1391:       TSBEULER - ODE solver using the implicit backward Euler method

1393:   Level: beginner

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

1398: .seealso: [](ch_ts), `TSCreate()`, `TS`, `TSSetType()`, `TSEULER`, `TSCN`, `TSTHETA`
1399: M*/
1400: PETSC_EXTERN PetscErrorCode TSCreate_BEuler(TS ts)
1401: {
1402:   PetscFunctionBegin;
1403:   PetscCall(TSCreate_Theta(ts));
1404:   PetscCall(TSThetaSetTheta(ts, 1.0));
1405:   PetscCall(TSThetaSetEndpoint(ts, PETSC_FALSE));
1406:   ts->ops->setup = TSSetUp_BEuler;
1407:   ts->ops->view  = TSView_BEuler;
1408:   PetscFunctionReturn(PETSC_SUCCESS);
1409: }

1411: static PetscErrorCode TSSetUp_CN(TS ts)
1412: {
1413:   TS_Theta *th = (TS_Theta *)ts->data;

1415:   PetscFunctionBegin;
1416:   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");
1417:   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");
1418:   PetscCall(TSSetUp_Theta(ts));
1419:   PetscFunctionReturn(PETSC_SUCCESS);
1420: }

1422: static PetscErrorCode TSView_CN(TS ts, PetscViewer viewer)
1423: {
1424:   PetscFunctionBegin;
1425:   PetscFunctionReturn(PETSC_SUCCESS);
1426: }

1428: /*MC
1429:       TSCN - ODE solver using the implicit Crank-Nicolson method.

1431:   Level: beginner

1433:   Notes:
1434:   `TSCN` is equivalent to `TSTHETA` with Theta=0.5 and the "endpoint" option set. I.e.
1435: .vb
1436:   -ts_type theta
1437:   -ts_theta_theta 0.5
1438:   -ts_theta_endpoint
1439: .ve

1441: .seealso: [](ch_ts), `TSCreate()`, `TS`, `TSSetType()`, `TSBEULER`, `TSTHETA`, `TSType`,
1442: M*/
1443: PETSC_EXTERN PetscErrorCode TSCreate_CN(TS ts)
1444: {
1445:   PetscFunctionBegin;
1446:   PetscCall(TSCreate_Theta(ts));
1447:   PetscCall(TSThetaSetTheta(ts, 0.5));
1448:   PetscCall(TSThetaSetEndpoint(ts, PETSC_TRUE));
1449:   ts->ops->setup = TSSetUp_CN;
1450:   ts->ops->view  = TSView_CN;
1451:   PetscFunctionReturn(PETSC_SUCCESS);
1452: }