Actual source code: theta.c

  1: /*
  2:   Code for timestepping with implicit Theta method
  3: */
  4: #include <petsc/private/tsimpl.h>
  5: #include <petscsnes.h>
  6: #include <petscdm.h>
  7: #include <petscmat.h>

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

722:   PetscFunctionBegin;
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_DETERMINE, &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_DETERMINE, &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_DETERMINE, &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_DETERMINE, &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_DETERMINE, &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:     if (!th->vec_sol_prev) PetscCall(VecDuplicate(ts->vec_sol, &th->vec_sol_prev));
1044:     if (!th->vec_lte_work) 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;

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

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

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

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

1184:    Level: beginner

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

1263:   ts->usessnes = PETSC_TRUE;

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

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

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

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

1286:   Not Collective

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

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

1294:   Level: advanced

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

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

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

1313:   Not Collective

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

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

1322:   Level: intermediate

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

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

1337:   Not Collective

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

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

1345:   Level: advanced

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

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

1361:   Not Collective

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

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

1370:   Level: intermediate

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

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

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

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

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

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

1407:   Level: beginner

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

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

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

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

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

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

1445:   Level: beginner

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

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