Actual source code: theta.c
1: /*
2: Code for timestepping with implicit Theta method
3: */
4: #include <petsc/private/tsimpl.h>
5: #include <petscsnes.h>
6: #include <petscdm.h>
7: #include <petscmat.h>
9: typedef struct {
10: /* context for time stepping */
11: PetscReal stage_time;
12: Vec Stages[2]; /* Storage for stage solutions */
13: Vec X0, X, Xdot; /* Storage for u^n, u^n + dt a_{11} k_1, and time derivative u^{n+1}_t */
14: Vec affine; /* Affine vector needed for residual at beginning of step in endpoint formulation */
15: PetscReal Theta;
16: PetscReal shift; /* Shift parameter for SNES Jacobian, used by forward, TLM and adjoint */
17: PetscInt order;
18: PetscBool endpoint;
19: PetscBool extrapolate;
20: TSStepStatus status;
21: Vec VecCostIntegral0; /* Backup for roll-backs due to events, used by cost integral */
22: PetscReal ptime0; /* Backup for ts->ptime, the start time of current time step, used by TLM and cost integral */
23: PetscReal time_step0; /* Backup for ts->timestep, the step size of current time step, used by TLM and cost integral*/
25: /* context for sensitivity analysis */
26: PetscInt num_tlm; /* Total number of tangent linear equations */
27: Vec *VecsDeltaLam; /* Increment of the adjoint sensitivity w.r.t IC at stage */
28: Vec *VecsDeltaMu; /* Increment of the adjoint sensitivity w.r.t P at stage */
29: Vec *VecsSensiTemp; /* Vector to be multiplied with Jacobian transpose */
30: Mat MatFwdStages[2]; /* TLM Stages */
31: Mat MatDeltaFwdSensip; /* Increment of the forward sensitivity at stage */
32: Vec VecDeltaFwdSensipCol; /* Working vector for holding one column of the sensitivity matrix */
33: Mat MatFwdSensip0; /* backup for roll-backs due to events */
34: Mat MatIntegralSensipTemp; /* Working vector for forward integral sensitivity */
35: Mat MatIntegralSensip0; /* backup for roll-backs due to events */
36: Vec *VecsDeltaLam2; /* Increment of the 2nd-order adjoint sensitivity w.r.t IC at stage */
37: Vec *VecsDeltaMu2; /* Increment of the 2nd-order adjoint sensitivity w.r.t P at stage */
38: Vec *VecsSensi2Temp; /* Working vectors that holds the residual for the second-order adjoint */
39: Vec *VecsAffine; /* Working vectors to store residuals */
40: /* context for error estimation */
41: Vec vec_sol_prev;
42: Vec vec_lte_work;
43: } TS_Theta;
45: static PetscErrorCode TSThetaGetX0AndXdot(TS ts, DM dm, Vec *X0, Vec *Xdot)
46: {
47: TS_Theta *th = (TS_Theta *)ts->data;
49: PetscFunctionBegin;
50: if (X0) {
51: if (dm && dm != ts->dm) {
52: PetscCall(DMGetNamedGlobalVector(dm, "TSTheta_X0", X0));
53: } else *X0 = ts->vec_sol;
54: }
55: if (Xdot) {
56: if (dm && dm != ts->dm) {
57: PetscCall(DMGetNamedGlobalVector(dm, "TSTheta_Xdot", Xdot));
58: } else *Xdot = th->Xdot;
59: }
60: PetscFunctionReturn(PETSC_SUCCESS);
61: }
63: static PetscErrorCode TSThetaRestoreX0AndXdot(TS ts, DM dm, Vec *X0, Vec *Xdot)
64: {
65: PetscFunctionBegin;
66: if (X0) {
67: if (dm && dm != ts->dm) PetscCall(DMRestoreNamedGlobalVector(dm, "TSTheta_X0", X0));
68: }
69: if (Xdot) {
70: if (dm && dm != ts->dm) PetscCall(DMRestoreNamedGlobalVector(dm, "TSTheta_Xdot", Xdot));
71: }
72: PetscFunctionReturn(PETSC_SUCCESS);
73: }
75: static PetscErrorCode DMCoarsenHook_TSTheta(DM fine, DM coarse, void *ctx)
76: {
77: PetscFunctionBegin;
78: PetscFunctionReturn(PETSC_SUCCESS);
79: }
81: static PetscErrorCode DMRestrictHook_TSTheta(DM fine, Mat restrct, Vec rscale, Mat inject, DM coarse, void *ctx)
82: {
83: TS ts = (TS)ctx;
84: Vec X0, Xdot, X0_c, Xdot_c;
86: PetscFunctionBegin;
87: PetscCall(TSThetaGetX0AndXdot(ts, fine, &X0, &Xdot));
88: PetscCall(TSThetaGetX0AndXdot(ts, coarse, &X0_c, &Xdot_c));
89: PetscCall(MatRestrict(restrct, X0, X0_c));
90: PetscCall(MatRestrict(restrct, Xdot, Xdot_c));
91: PetscCall(VecPointwiseMult(X0_c, rscale, X0_c));
92: PetscCall(VecPointwiseMult(Xdot_c, rscale, Xdot_c));
93: PetscCall(TSThetaRestoreX0AndXdot(ts, fine, &X0, &Xdot));
94: PetscCall(TSThetaRestoreX0AndXdot(ts, coarse, &X0_c, &Xdot_c));
95: PetscFunctionReturn(PETSC_SUCCESS);
96: }
98: static PetscErrorCode DMSubDomainHook_TSTheta(DM dm, DM subdm, void *ctx)
99: {
100: PetscFunctionBegin;
101: PetscFunctionReturn(PETSC_SUCCESS);
102: }
104: static PetscErrorCode DMSubDomainRestrictHook_TSTheta(DM dm, VecScatter gscat, VecScatter lscat, DM subdm, void *ctx)
105: {
106: TS ts = (TS)ctx;
107: Vec X0, Xdot, X0_sub, Xdot_sub;
109: PetscFunctionBegin;
110: PetscCall(TSThetaGetX0AndXdot(ts, dm, &X0, &Xdot));
111: PetscCall(TSThetaGetX0AndXdot(ts, subdm, &X0_sub, &Xdot_sub));
113: PetscCall(VecScatterBegin(gscat, X0, X0_sub, INSERT_VALUES, SCATTER_FORWARD));
114: PetscCall(VecScatterEnd(gscat, X0, X0_sub, INSERT_VALUES, SCATTER_FORWARD));
116: PetscCall(VecScatterBegin(gscat, Xdot, Xdot_sub, INSERT_VALUES, SCATTER_FORWARD));
117: PetscCall(VecScatterEnd(gscat, Xdot, Xdot_sub, INSERT_VALUES, SCATTER_FORWARD));
119: PetscCall(TSThetaRestoreX0AndXdot(ts, dm, &X0, &Xdot));
120: PetscCall(TSThetaRestoreX0AndXdot(ts, subdm, &X0_sub, &Xdot_sub));
121: PetscFunctionReturn(PETSC_SUCCESS);
122: }
124: static PetscErrorCode TSThetaEvaluateCostIntegral(TS ts)
125: {
126: TS_Theta *th = (TS_Theta *)ts->data;
127: TS quadts = ts->quadraturets;
129: PetscFunctionBegin;
130: if (th->endpoint) {
131: /* Evolve ts->vec_costintegral to compute integrals */
132: if (th->Theta != 1.0) {
133: PetscCall(TSComputeRHSFunction(quadts, th->ptime0, th->X0, ts->vec_costintegrand));
134: PetscCall(VecAXPY(quadts->vec_sol, th->time_step0 * (1.0 - th->Theta), ts->vec_costintegrand));
135: }
136: PetscCall(TSComputeRHSFunction(quadts, ts->ptime, ts->vec_sol, ts->vec_costintegrand));
137: PetscCall(VecAXPY(quadts->vec_sol, th->time_step0 * th->Theta, ts->vec_costintegrand));
138: } else {
139: PetscCall(TSComputeRHSFunction(quadts, th->stage_time, th->X, ts->vec_costintegrand));
140: PetscCall(VecAXPY(quadts->vec_sol, th->time_step0, ts->vec_costintegrand));
141: }
142: PetscFunctionReturn(PETSC_SUCCESS);
143: }
145: static PetscErrorCode TSForwardCostIntegral_Theta(TS ts)
146: {
147: TS_Theta *th = (TS_Theta *)ts->data;
148: TS quadts = ts->quadraturets;
150: PetscFunctionBegin;
151: /* backup cost integral */
152: PetscCall(VecCopy(quadts->vec_sol, th->VecCostIntegral0));
153: PetscCall(TSThetaEvaluateCostIntegral(ts));
154: PetscFunctionReturn(PETSC_SUCCESS);
155: }
157: static PetscErrorCode TSAdjointCostIntegral_Theta(TS ts)
158: {
159: TS_Theta *th = (TS_Theta *)ts->data;
161: PetscFunctionBegin;
162: /* Like TSForwardCostIntegral(), the adjoint cost integral evaluation relies on ptime0 and time_step0. */
163: th->ptime0 = ts->ptime + ts->time_step;
164: th->time_step0 = -ts->time_step;
165: PetscCall(TSThetaEvaluateCostIntegral(ts));
166: PetscFunctionReturn(PETSC_SUCCESS);
167: }
169: static PetscErrorCode TSTheta_SNESSolve(TS ts, Vec b, Vec x)
170: {
171: PetscInt nits, lits;
173: PetscFunctionBegin;
174: PetscCall(SNESSolve(ts->snes, b, x));
175: PetscCall(SNESGetIterationNumber(ts->snes, &nits));
176: PetscCall(SNESGetLinearSolveIterations(ts->snes, &lits));
177: ts->snes_its += nits;
178: ts->ksp_its += lits;
179: PetscFunctionReturn(PETSC_SUCCESS);
180: }
182: /* We need to transfer X0 which will be copied into sol_prev */
183: static PetscErrorCode TSResizeRegister_Theta(TS ts, PetscBool reg)
184: {
185: TS_Theta *th = (TS_Theta *)ts->data;
186: const char name[] = "ts:theta:X0";
188: PetscFunctionBegin;
189: if (reg && th->vec_sol_prev) {
190: PetscCall(TSResizeRegisterVec(ts, name, th->X0));
191: } else if (!reg) {
192: PetscCall(TSResizeRetrieveVec(ts, name, &th->X0));
193: PetscCall(PetscObjectReference((PetscObject)th->X0));
194: }
195: PetscFunctionReturn(PETSC_SUCCESS);
196: }
198: static PetscErrorCode TSStep_Theta(TS ts)
199: {
200: TS_Theta *th = (TS_Theta *)ts->data;
201: PetscInt rejections = 0;
202: PetscBool stageok, accept = PETSC_TRUE;
203: PetscReal next_time_step = ts->time_step;
205: PetscFunctionBegin;
206: if (!ts->steprollback) {
207: if (th->vec_sol_prev) PetscCall(VecCopy(th->X0, th->vec_sol_prev));
208: PetscCall(VecCopy(ts->vec_sol, th->X0));
209: }
211: th->status = TS_STEP_INCOMPLETE;
212: while (!ts->reason && th->status != TS_STEP_COMPLETE) {
213: th->shift = 1 / (th->Theta * ts->time_step);
214: th->stage_time = ts->ptime + (th->endpoint ? (PetscReal)1 : th->Theta) * ts->time_step;
215: PetscCall(VecCopy(th->X0, th->X));
216: if (th->extrapolate && !ts->steprestart) PetscCall(VecAXPY(th->X, 1 / th->shift, th->Xdot));
217: if (th->endpoint) { /* This formulation assumes linear time-independent mass matrix */
218: if (!th->affine) PetscCall(VecDuplicate(ts->vec_sol, &th->affine));
219: PetscCall(VecZeroEntries(th->Xdot));
220: PetscCall(TSComputeIFunction(ts, ts->ptime, th->X0, th->Xdot, th->affine, PETSC_FALSE));
221: PetscCall(VecScale(th->affine, (th->Theta - 1) / th->Theta));
222: }
223: PetscCall(TSPreStage(ts, th->stage_time));
224: PetscCall(TSTheta_SNESSolve(ts, th->endpoint ? th->affine : NULL, th->X));
225: PetscCall(TSPostStage(ts, th->stage_time, 0, &th->X));
226: PetscCall(TSAdaptCheckStage(ts->adapt, ts, th->stage_time, th->X, &stageok));
227: if (!stageok) goto reject_step;
229: th->status = TS_STEP_PENDING;
230: if (th->endpoint) {
231: PetscCall(VecCopy(th->X, ts->vec_sol));
232: } else {
233: PetscCall(VecAXPBYPCZ(th->Xdot, -th->shift, th->shift, 0, th->X0, th->X));
234: PetscCall(VecAXPY(ts->vec_sol, ts->time_step, th->Xdot));
235: }
236: PetscCall(TSAdaptChoose(ts->adapt, ts, ts->time_step, NULL, &next_time_step, &accept));
237: th->status = accept ? TS_STEP_COMPLETE : TS_STEP_INCOMPLETE;
238: if (!accept) {
239: PetscCall(VecCopy(th->X0, ts->vec_sol));
240: ts->time_step = next_time_step;
241: goto reject_step;
242: }
244: if (ts->forward_solve || ts->costintegralfwd) { /* Save the info for the later use in cost integral evaluation */
245: th->ptime0 = ts->ptime;
246: th->time_step0 = ts->time_step;
247: }
248: ts->ptime += ts->time_step;
249: ts->time_step = next_time_step;
250: break;
252: reject_step:
253: ts->reject++;
254: accept = PETSC_FALSE;
255: if (!ts->reason && ++rejections > ts->max_reject && ts->max_reject >= 0) {
256: ts->reason = TS_DIVERGED_STEP_REJECTED;
257: PetscCall(PetscInfo(ts, "Step=%" PetscInt_FMT ", step rejections %" PetscInt_FMT " greater than current TS allowed, stopping solve\n", ts->steps, rejections));
258: }
259: }
260: PetscFunctionReturn(PETSC_SUCCESS);
261: }
263: static PetscErrorCode TSAdjointStepBEuler_Private(TS ts)
264: {
265: TS_Theta *th = (TS_Theta *)ts->data;
266: TS quadts = ts->quadraturets;
267: Vec *VecsDeltaLam = th->VecsDeltaLam, *VecsDeltaMu = th->VecsDeltaMu, *VecsSensiTemp = th->VecsSensiTemp;
268: Vec *VecsDeltaLam2 = th->VecsDeltaLam2, *VecsDeltaMu2 = th->VecsDeltaMu2, *VecsSensi2Temp = th->VecsSensi2Temp;
269: PetscInt nadj;
270: Mat J, Jpre, quadJ = NULL, quadJp = NULL;
271: KSP ksp;
272: PetscScalar *xarr;
273: TSEquationType eqtype;
274: PetscBool isexplicitode = PETSC_FALSE;
275: PetscReal adjoint_time_step;
277: PetscFunctionBegin;
278: PetscCall(TSGetEquationType(ts, &eqtype));
279: if (eqtype == TS_EQ_ODE_EXPLICIT) {
280: isexplicitode = PETSC_TRUE;
281: VecsDeltaLam = ts->vecs_sensi;
282: VecsDeltaLam2 = ts->vecs_sensi2;
283: }
284: th->status = TS_STEP_INCOMPLETE;
285: PetscCall(SNESGetKSP(ts->snes, &ksp));
286: PetscCall(TSGetIJacobian(ts, &J, &Jpre, NULL, NULL));
287: if (quadts) {
288: PetscCall(TSGetRHSJacobian(quadts, &quadJ, NULL, NULL, NULL));
289: PetscCall(TSGetRHSJacobianP(quadts, &quadJp, NULL, NULL));
290: }
292: th->stage_time = ts->ptime;
293: adjoint_time_step = -ts->time_step; /* always positive since time_step is negative */
295: /* Build RHS for first-order adjoint lambda_{n+1}/h + r_u^T(n+1) */
296: if (quadts) PetscCall(TSComputeRHSJacobian(quadts, th->stage_time, ts->vec_sol, quadJ, NULL));
298: for (nadj = 0; nadj < ts->numcost; nadj++) {
299: PetscCall(VecCopy(ts->vecs_sensi[nadj], VecsSensiTemp[nadj]));
300: PetscCall(VecScale(VecsSensiTemp[nadj], 1. / adjoint_time_step)); /* lambda_{n+1}/h */
301: if (quadJ) {
302: PetscCall(MatDenseGetColumn(quadJ, nadj, &xarr));
303: PetscCall(VecPlaceArray(ts->vec_drdu_col, xarr));
304: PetscCall(VecAXPY(VecsSensiTemp[nadj], 1., ts->vec_drdu_col));
305: PetscCall(VecResetArray(ts->vec_drdu_col));
306: PetscCall(MatDenseRestoreColumn(quadJ, &xarr));
307: }
308: }
310: /* Build LHS for first-order adjoint */
311: th->shift = 1. / adjoint_time_step;
312: PetscCall(TSComputeSNESJacobian(ts, ts->vec_sol, J, Jpre));
313: PetscCall(KSPSetOperators(ksp, J, Jpre));
315: /* Solve stage equation LHS*lambda_s = RHS for first-order adjoint */
316: for (nadj = 0; nadj < ts->numcost; nadj++) {
317: KSPConvergedReason kspreason;
318: PetscCall(KSPSolveTranspose(ksp, VecsSensiTemp[nadj], VecsDeltaLam[nadj]));
319: PetscCall(KSPGetConvergedReason(ksp, &kspreason));
320: if (kspreason < 0) {
321: ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
322: PetscCall(PetscInfo(ts, "Step=%" PetscInt_FMT ", %" PetscInt_FMT "th cost function, transposed linear solve fails, stopping 1st-order adjoint solve\n", ts->steps, nadj));
323: }
324: }
326: if (ts->vecs_sensi2) { /* U_{n+1} */
327: /* Get w1 at t_{n+1} from TLM matrix */
328: PetscCall(MatDenseGetColumn(ts->mat_sensip, 0, &xarr));
329: PetscCall(VecPlaceArray(ts->vec_sensip_col, xarr));
330: /* lambda_s^T F_UU w_1 */
331: PetscCall(TSComputeIHessianProductFunctionUU(ts, th->stage_time, ts->vec_sol, VecsDeltaLam, ts->vec_sensip_col, ts->vecs_fuu));
332: /* lambda_s^T F_UP w_2 */
333: PetscCall(TSComputeIHessianProductFunctionUP(ts, th->stage_time, ts->vec_sol, VecsDeltaLam, ts->vec_dir, ts->vecs_fup));
334: for (nadj = 0; nadj < ts->numcost; nadj++) { /* compute the residual */
335: PetscCall(VecCopy(ts->vecs_sensi2[nadj], VecsSensi2Temp[nadj]));
336: PetscCall(VecScale(VecsSensi2Temp[nadj], 1. / adjoint_time_step));
337: PetscCall(VecAXPY(VecsSensi2Temp[nadj], -1., ts->vecs_fuu[nadj]));
338: if (ts->vecs_fup) PetscCall(VecAXPY(VecsSensi2Temp[nadj], -1., ts->vecs_fup[nadj]));
339: }
340: /* Solve stage equation LHS X = RHS for second-order adjoint */
341: for (nadj = 0; nadj < ts->numcost; nadj++) {
342: KSPConvergedReason kspreason;
343: PetscCall(KSPSolveTranspose(ksp, VecsSensi2Temp[nadj], VecsDeltaLam2[nadj]));
344: PetscCall(KSPGetConvergedReason(ksp, &kspreason));
345: if (kspreason < 0) {
346: ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
347: PetscCall(PetscInfo(ts, "Step=%" PetscInt_FMT ", %" PetscInt_FMT "th cost function, transposed linear solve fails, stopping 2nd-order adjoint solve\n", ts->steps, nadj));
348: }
349: }
350: }
352: /* Update sensitivities, and evaluate integrals if there is any */
353: if (!isexplicitode) {
354: th->shift = 0.0;
355: PetscCall(TSComputeSNESJacobian(ts, ts->vec_sol, J, Jpre));
356: PetscCall(KSPSetOperators(ksp, J, Jpre));
357: for (nadj = 0; nadj < ts->numcost; nadj++) {
358: /* Add f_U \lambda_s to the original RHS */
359: PetscCall(VecScale(VecsSensiTemp[nadj], -1.));
360: PetscCall(MatMultTransposeAdd(J, VecsDeltaLam[nadj], VecsSensiTemp[nadj], VecsSensiTemp[nadj]));
361: PetscCall(VecScale(VecsSensiTemp[nadj], -adjoint_time_step));
362: PetscCall(VecCopy(VecsSensiTemp[nadj], ts->vecs_sensi[nadj]));
363: if (ts->vecs_sensi2) {
364: PetscCall(MatMultTransposeAdd(J, VecsDeltaLam2[nadj], VecsSensi2Temp[nadj], VecsSensi2Temp[nadj]));
365: PetscCall(VecScale(VecsSensi2Temp[nadj], -adjoint_time_step));
366: PetscCall(VecCopy(VecsSensi2Temp[nadj], ts->vecs_sensi2[nadj]));
367: }
368: }
369: }
370: if (ts->vecs_sensip) {
371: PetscCall(VecAXPBYPCZ(th->Xdot, -1. / adjoint_time_step, 1.0 / adjoint_time_step, 0, th->X0, ts->vec_sol));
372: PetscCall(TSComputeIJacobianP(ts, th->stage_time, ts->vec_sol, th->Xdot, 1. / adjoint_time_step, ts->Jacp, PETSC_FALSE)); /* get -f_p */
373: if (quadts) PetscCall(TSComputeRHSJacobianP(quadts, th->stage_time, ts->vec_sol, quadJp));
374: if (ts->vecs_sensi2p) {
375: /* lambda_s^T F_PU w_1 */
376: PetscCall(TSComputeIHessianProductFunctionPU(ts, th->stage_time, ts->vec_sol, VecsDeltaLam, ts->vec_sensip_col, ts->vecs_fpu));
377: /* lambda_s^T F_PP w_2 */
378: PetscCall(TSComputeIHessianProductFunctionPP(ts, th->stage_time, ts->vec_sol, VecsDeltaLam, ts->vec_dir, ts->vecs_fpp));
379: }
381: for (nadj = 0; nadj < ts->numcost; nadj++) {
382: PetscCall(MatMultTranspose(ts->Jacp, VecsDeltaLam[nadj], VecsDeltaMu[nadj]));
383: PetscCall(VecAXPY(ts->vecs_sensip[nadj], -adjoint_time_step, VecsDeltaMu[nadj]));
384: if (quadJp) {
385: PetscCall(MatDenseGetColumn(quadJp, nadj, &xarr));
386: PetscCall(VecPlaceArray(ts->vec_drdp_col, xarr));
387: PetscCall(VecAXPY(ts->vecs_sensip[nadj], adjoint_time_step, ts->vec_drdp_col));
388: PetscCall(VecResetArray(ts->vec_drdp_col));
389: PetscCall(MatDenseRestoreColumn(quadJp, &xarr));
390: }
391: if (ts->vecs_sensi2p) {
392: PetscCall(MatMultTranspose(ts->Jacp, VecsDeltaLam2[nadj], VecsDeltaMu2[nadj]));
393: PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step, VecsDeltaMu2[nadj]));
394: if (ts->vecs_fpu) PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step, ts->vecs_fpu[nadj]));
395: if (ts->vecs_fpp) PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step, ts->vecs_fpp[nadj]));
396: }
397: }
398: }
400: if (ts->vecs_sensi2) {
401: PetscCall(VecResetArray(ts->vec_sensip_col));
402: PetscCall(MatDenseRestoreColumn(ts->mat_sensip, &xarr));
403: }
404: th->status = TS_STEP_COMPLETE;
405: PetscFunctionReturn(PETSC_SUCCESS);
406: }
408: static PetscErrorCode TSAdjointStep_Theta(TS ts)
409: {
410: TS_Theta *th = (TS_Theta *)ts->data;
411: TS quadts = ts->quadraturets;
412: Vec *VecsDeltaLam = th->VecsDeltaLam, *VecsDeltaMu = th->VecsDeltaMu, *VecsSensiTemp = th->VecsSensiTemp;
413: Vec *VecsDeltaLam2 = th->VecsDeltaLam2, *VecsDeltaMu2 = th->VecsDeltaMu2, *VecsSensi2Temp = th->VecsSensi2Temp;
414: PetscInt nadj;
415: Mat J, Jpre, quadJ = NULL, quadJp = NULL;
416: KSP ksp;
417: PetscScalar *xarr;
418: PetscReal adjoint_time_step;
419: PetscReal adjoint_ptime; /* end time of the adjoint time step (ts->ptime is the start time, usually ts->ptime is larger than adjoint_ptime) */
421: PetscFunctionBegin;
422: if (th->Theta == 1.) {
423: PetscCall(TSAdjointStepBEuler_Private(ts));
424: PetscFunctionReturn(PETSC_SUCCESS);
425: }
426: th->status = TS_STEP_INCOMPLETE;
427: PetscCall(SNESGetKSP(ts->snes, &ksp));
428: PetscCall(TSGetIJacobian(ts, &J, &Jpre, NULL, NULL));
429: if (quadts) {
430: PetscCall(TSGetRHSJacobian(quadts, &quadJ, NULL, NULL, NULL));
431: PetscCall(TSGetRHSJacobianP(quadts, &quadJp, NULL, NULL));
432: }
433: /* If endpoint=1, th->ptime and th->X0 will be used; if endpoint=0, th->stage_time and th->X will be used. */
434: th->stage_time = th->endpoint ? ts->ptime : (ts->ptime + (1. - th->Theta) * ts->time_step);
435: adjoint_ptime = ts->ptime + ts->time_step;
436: adjoint_time_step = -ts->time_step; /* always positive since time_step is negative */
438: if (!th->endpoint) {
439: /* recover th->X0 using vec_sol and the stage value th->X */
440: PetscCall(VecAXPBYPCZ(th->X0, 1.0 / (1.0 - th->Theta), th->Theta / (th->Theta - 1.0), 0, th->X, ts->vec_sol));
441: }
443: /* Build RHS for first-order adjoint */
444: /* Cost function has an integral term */
445: if (quadts) {
446: if (th->endpoint) {
447: PetscCall(TSComputeRHSJacobian(quadts, th->stage_time, ts->vec_sol, quadJ, NULL));
448: } else {
449: PetscCall(TSComputeRHSJacobian(quadts, th->stage_time, th->X, quadJ, NULL));
450: }
451: }
453: for (nadj = 0; nadj < ts->numcost; nadj++) {
454: PetscCall(VecCopy(ts->vecs_sensi[nadj], VecsSensiTemp[nadj]));
455: PetscCall(VecScale(VecsSensiTemp[nadj], 1. / (th->Theta * adjoint_time_step)));
456: if (quadJ) {
457: PetscCall(MatDenseGetColumn(quadJ, nadj, &xarr));
458: PetscCall(VecPlaceArray(ts->vec_drdu_col, xarr));
459: PetscCall(VecAXPY(VecsSensiTemp[nadj], 1., ts->vec_drdu_col));
460: PetscCall(VecResetArray(ts->vec_drdu_col));
461: PetscCall(MatDenseRestoreColumn(quadJ, &xarr));
462: }
463: }
465: /* Build LHS for first-order adjoint */
466: th->shift = 1. / (th->Theta * adjoint_time_step);
467: if (th->endpoint) {
468: PetscCall(TSComputeSNESJacobian(ts, ts->vec_sol, J, Jpre));
469: } else {
470: PetscCall(TSComputeSNESJacobian(ts, th->X, J, Jpre));
471: }
472: PetscCall(KSPSetOperators(ksp, J, Jpre));
474: /* Solve stage equation LHS*lambda_s = RHS for first-order adjoint */
475: for (nadj = 0; nadj < ts->numcost; nadj++) {
476: KSPConvergedReason kspreason;
477: PetscCall(KSPSolveTranspose(ksp, VecsSensiTemp[nadj], VecsDeltaLam[nadj]));
478: PetscCall(KSPGetConvergedReason(ksp, &kspreason));
479: if (kspreason < 0) {
480: ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
481: PetscCall(PetscInfo(ts, "Step=%" PetscInt_FMT ", %" PetscInt_FMT "th cost function, transposed linear solve fails, stopping 1st-order adjoint solve\n", ts->steps, nadj));
482: }
483: }
485: /* Second-order adjoint */
486: if (ts->vecs_sensi2) { /* U_{n+1} */
487: PetscCheck(th->endpoint, PetscObjectComm((PetscObject)ts), PETSC_ERR_SUP, "Operation not implemented in TS_Theta");
488: /* Get w1 at t_{n+1} from TLM matrix */
489: PetscCall(MatDenseGetColumn(ts->mat_sensip, 0, &xarr));
490: PetscCall(VecPlaceArray(ts->vec_sensip_col, xarr));
491: /* lambda_s^T F_UU w_1 */
492: PetscCall(TSComputeIHessianProductFunctionUU(ts, th->stage_time, ts->vec_sol, VecsDeltaLam, ts->vec_sensip_col, ts->vecs_fuu));
493: PetscCall(VecResetArray(ts->vec_sensip_col));
494: PetscCall(MatDenseRestoreColumn(ts->mat_sensip, &xarr));
495: /* lambda_s^T F_UP w_2 */
496: PetscCall(TSComputeIHessianProductFunctionUP(ts, th->stage_time, ts->vec_sol, VecsDeltaLam, ts->vec_dir, ts->vecs_fup));
497: for (nadj = 0; nadj < ts->numcost; nadj++) { /* compute the residual */
498: PetscCall(VecCopy(ts->vecs_sensi2[nadj], VecsSensi2Temp[nadj]));
499: PetscCall(VecScale(VecsSensi2Temp[nadj], th->shift));
500: PetscCall(VecAXPY(VecsSensi2Temp[nadj], -1., ts->vecs_fuu[nadj]));
501: if (ts->vecs_fup) PetscCall(VecAXPY(VecsSensi2Temp[nadj], -1., ts->vecs_fup[nadj]));
502: }
503: /* Solve stage equation LHS X = RHS for second-order adjoint */
504: for (nadj = 0; nadj < ts->numcost; nadj++) {
505: KSPConvergedReason kspreason;
506: PetscCall(KSPSolveTranspose(ksp, VecsSensi2Temp[nadj], VecsDeltaLam2[nadj]));
507: PetscCall(KSPGetConvergedReason(ksp, &kspreason));
508: if (kspreason < 0) {
509: ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
510: PetscCall(PetscInfo(ts, "Step=%" PetscInt_FMT ", %" PetscInt_FMT "th cost function, transposed linear solve fails, stopping 2nd-order adjoint solve\n", ts->steps, nadj));
511: }
512: }
513: }
515: /* Update sensitivities, and evaluate integrals if there is any */
516: if (th->endpoint) { /* two-stage Theta methods with th->Theta!=1, th->Theta==1 leads to BEuler */
517: th->shift = 1. / ((th->Theta - 1.) * adjoint_time_step);
518: th->stage_time = adjoint_ptime;
519: PetscCall(TSComputeSNESJacobian(ts, th->X0, J, Jpre));
520: PetscCall(KSPSetOperators(ksp, J, Jpre));
521: /* R_U at t_n */
522: if (quadts) PetscCall(TSComputeRHSJacobian(quadts, adjoint_ptime, th->X0, quadJ, NULL));
523: for (nadj = 0; nadj < ts->numcost; nadj++) {
524: PetscCall(MatMultTranspose(J, VecsDeltaLam[nadj], ts->vecs_sensi[nadj]));
525: if (quadJ) {
526: PetscCall(MatDenseGetColumn(quadJ, nadj, &xarr));
527: PetscCall(VecPlaceArray(ts->vec_drdu_col, xarr));
528: PetscCall(VecAXPY(ts->vecs_sensi[nadj], -1., ts->vec_drdu_col));
529: PetscCall(VecResetArray(ts->vec_drdu_col));
530: PetscCall(MatDenseRestoreColumn(quadJ, &xarr));
531: }
532: PetscCall(VecScale(ts->vecs_sensi[nadj], 1. / th->shift));
533: }
535: /* Second-order adjoint */
536: if (ts->vecs_sensi2) { /* U_n */
537: /* Get w1 at t_n from TLM matrix */
538: PetscCall(MatDenseGetColumn(th->MatFwdSensip0, 0, &xarr));
539: PetscCall(VecPlaceArray(ts->vec_sensip_col, xarr));
540: /* lambda_s^T F_UU w_1 */
541: PetscCall(TSComputeIHessianProductFunctionUU(ts, adjoint_ptime, th->X0, VecsDeltaLam, ts->vec_sensip_col, ts->vecs_fuu));
542: PetscCall(VecResetArray(ts->vec_sensip_col));
543: PetscCall(MatDenseRestoreColumn(th->MatFwdSensip0, &xarr));
544: /* lambda_s^T F_UU w_2 */
545: PetscCall(TSComputeIHessianProductFunctionUP(ts, adjoint_ptime, th->X0, VecsDeltaLam, ts->vec_dir, ts->vecs_fup));
546: for (nadj = 0; nadj < ts->numcost; nadj++) {
547: /* M^T Lambda_s + h(1-theta) F_U^T Lambda_s + h(1-theta) lambda_s^T F_UU w_1 + lambda_s^T F_UP w_2 */
548: PetscCall(MatMultTranspose(J, VecsDeltaLam2[nadj], ts->vecs_sensi2[nadj]));
549: PetscCall(VecAXPY(ts->vecs_sensi2[nadj], 1., ts->vecs_fuu[nadj]));
550: if (ts->vecs_fup) PetscCall(VecAXPY(ts->vecs_sensi2[nadj], 1., ts->vecs_fup[nadj]));
551: PetscCall(VecScale(ts->vecs_sensi2[nadj], 1. / th->shift));
552: }
553: }
555: th->stage_time = ts->ptime; /* recover the old value */
557: if (ts->vecs_sensip) { /* sensitivities wrt parameters */
558: /* U_{n+1} */
559: th->shift = 1.0 / (adjoint_time_step * th->Theta);
560: PetscCall(VecAXPBYPCZ(th->Xdot, -th->shift, th->shift, 0, th->X0, ts->vec_sol));
561: PetscCall(TSComputeIJacobianP(ts, th->stage_time, ts->vec_sol, th->Xdot, -1. / (th->Theta * adjoint_time_step), ts->Jacp, PETSC_FALSE));
562: if (quadts) PetscCall(TSComputeRHSJacobianP(quadts, th->stage_time, ts->vec_sol, quadJp));
563: for (nadj = 0; nadj < ts->numcost; nadj++) {
564: PetscCall(MatMultTranspose(ts->Jacp, VecsDeltaLam[nadj], VecsDeltaMu[nadj]));
565: PetscCall(VecAXPY(ts->vecs_sensip[nadj], -adjoint_time_step * th->Theta, VecsDeltaMu[nadj]));
566: if (quadJp) {
567: PetscCall(MatDenseGetColumn(quadJp, nadj, &xarr));
568: PetscCall(VecPlaceArray(ts->vec_drdp_col, xarr));
569: PetscCall(VecAXPY(ts->vecs_sensip[nadj], adjoint_time_step * th->Theta, ts->vec_drdp_col));
570: PetscCall(VecResetArray(ts->vec_drdp_col));
571: PetscCall(MatDenseRestoreColumn(quadJp, &xarr));
572: }
573: }
574: if (ts->vecs_sensi2p) { /* second-order */
575: /* Get w1 at t_{n+1} from TLM matrix */
576: PetscCall(MatDenseGetColumn(ts->mat_sensip, 0, &xarr));
577: PetscCall(VecPlaceArray(ts->vec_sensip_col, xarr));
578: /* lambda_s^T F_PU w_1 */
579: PetscCall(TSComputeIHessianProductFunctionPU(ts, th->stage_time, ts->vec_sol, VecsDeltaLam, ts->vec_sensip_col, ts->vecs_fpu));
580: PetscCall(VecResetArray(ts->vec_sensip_col));
581: PetscCall(MatDenseRestoreColumn(ts->mat_sensip, &xarr));
583: /* lambda_s^T F_PP w_2 */
584: PetscCall(TSComputeIHessianProductFunctionPP(ts, th->stage_time, ts->vec_sol, VecsDeltaLam, ts->vec_dir, ts->vecs_fpp));
585: for (nadj = 0; nadj < ts->numcost; nadj++) {
586: /* Mu2 <- Mu2 + h theta F_P^T Lambda_s + h theta (lambda_s^T F_UU w_1 + lambda_s^T F_UP w_2) */
587: PetscCall(MatMultTranspose(ts->Jacp, VecsDeltaLam2[nadj], VecsDeltaMu2[nadj]));
588: PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step * th->Theta, VecsDeltaMu2[nadj]));
589: if (ts->vecs_fpu) PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step * th->Theta, ts->vecs_fpu[nadj]));
590: if (ts->vecs_fpp) PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step * th->Theta, ts->vecs_fpp[nadj]));
591: }
592: }
594: /* U_s */
595: PetscCall(VecZeroEntries(th->Xdot));
596: PetscCall(TSComputeIJacobianP(ts, adjoint_ptime, th->X0, th->Xdot, 1. / ((th->Theta - 1.0) * adjoint_time_step), ts->Jacp, PETSC_FALSE));
597: if (quadts) PetscCall(TSComputeRHSJacobianP(quadts, adjoint_ptime, th->X0, quadJp));
598: for (nadj = 0; nadj < ts->numcost; nadj++) {
599: PetscCall(MatMultTranspose(ts->Jacp, VecsDeltaLam[nadj], VecsDeltaMu[nadj]));
600: PetscCall(VecAXPY(ts->vecs_sensip[nadj], -adjoint_time_step * (1.0 - th->Theta), VecsDeltaMu[nadj]));
601: if (quadJp) {
602: PetscCall(MatDenseGetColumn(quadJp, nadj, &xarr));
603: PetscCall(VecPlaceArray(ts->vec_drdp_col, xarr));
604: PetscCall(VecAXPY(ts->vecs_sensip[nadj], adjoint_time_step * (1.0 - th->Theta), ts->vec_drdp_col));
605: PetscCall(VecResetArray(ts->vec_drdp_col));
606: PetscCall(MatDenseRestoreColumn(quadJp, &xarr));
607: }
608: if (ts->vecs_sensi2p) { /* second-order */
609: /* Get w1 at t_n from TLM matrix */
610: PetscCall(MatDenseGetColumn(th->MatFwdSensip0, 0, &xarr));
611: PetscCall(VecPlaceArray(ts->vec_sensip_col, xarr));
612: /* lambda_s^T F_PU w_1 */
613: PetscCall(TSComputeIHessianProductFunctionPU(ts, adjoint_ptime, th->X0, VecsDeltaLam, ts->vec_sensip_col, ts->vecs_fpu));
614: PetscCall(VecResetArray(ts->vec_sensip_col));
615: PetscCall(MatDenseRestoreColumn(th->MatFwdSensip0, &xarr));
616: /* lambda_s^T F_PP w_2 */
617: PetscCall(TSComputeIHessianProductFunctionPP(ts, adjoint_ptime, th->X0, VecsDeltaLam, ts->vec_dir, ts->vecs_fpp));
618: for (nadj = 0; nadj < ts->numcost; nadj++) {
619: /* Mu2 <- Mu2 + h(1-theta) F_P^T Lambda_s + h(1-theta) (lambda_s^T F_UU w_1 + lambda_s^T F_UP w_2) */
620: PetscCall(MatMultTranspose(ts->Jacp, VecsDeltaLam2[nadj], VecsDeltaMu2[nadj]));
621: PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step * (1.0 - th->Theta), VecsDeltaMu2[nadj]));
622: if (ts->vecs_fpu) PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step * (1.0 - th->Theta), ts->vecs_fpu[nadj]));
623: if (ts->vecs_fpp) PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step * (1.0 - th->Theta), ts->vecs_fpp[nadj]));
624: }
625: }
626: }
627: }
628: } else { /* one-stage case */
629: th->shift = 0.0;
630: PetscCall(TSComputeSNESJacobian(ts, th->X, J, Jpre)); /* get -f_y */
631: PetscCall(KSPSetOperators(ksp, J, Jpre));
632: if (quadts) PetscCall(TSComputeRHSJacobian(quadts, th->stage_time, th->X, quadJ, NULL));
633: for (nadj = 0; nadj < ts->numcost; nadj++) {
634: PetscCall(MatMultTranspose(J, VecsDeltaLam[nadj], VecsSensiTemp[nadj]));
635: PetscCall(VecAXPY(ts->vecs_sensi[nadj], -adjoint_time_step, VecsSensiTemp[nadj]));
636: if (quadJ) {
637: PetscCall(MatDenseGetColumn(quadJ, nadj, &xarr));
638: PetscCall(VecPlaceArray(ts->vec_drdu_col, xarr));
639: PetscCall(VecAXPY(ts->vecs_sensi[nadj], adjoint_time_step, ts->vec_drdu_col));
640: PetscCall(VecResetArray(ts->vec_drdu_col));
641: PetscCall(MatDenseRestoreColumn(quadJ, &xarr));
642: }
643: }
644: if (ts->vecs_sensip) {
645: th->shift = 1.0 / (adjoint_time_step * th->Theta);
646: PetscCall(VecAXPBYPCZ(th->Xdot, -th->shift, th->shift, 0, th->X0, th->X));
647: PetscCall(TSComputeIJacobianP(ts, th->stage_time, th->X, th->Xdot, th->shift, ts->Jacp, PETSC_FALSE));
648: if (quadts) PetscCall(TSComputeRHSJacobianP(quadts, th->stage_time, th->X, quadJp));
649: for (nadj = 0; nadj < ts->numcost; nadj++) {
650: PetscCall(MatMultTranspose(ts->Jacp, VecsDeltaLam[nadj], VecsDeltaMu[nadj]));
651: PetscCall(VecAXPY(ts->vecs_sensip[nadj], -adjoint_time_step, VecsDeltaMu[nadj]));
652: if (quadJp) {
653: PetscCall(MatDenseGetColumn(quadJp, nadj, &xarr));
654: PetscCall(VecPlaceArray(ts->vec_drdp_col, xarr));
655: PetscCall(VecAXPY(ts->vecs_sensip[nadj], adjoint_time_step, ts->vec_drdp_col));
656: PetscCall(VecResetArray(ts->vec_drdp_col));
657: PetscCall(MatDenseRestoreColumn(quadJp, &xarr));
658: }
659: }
660: }
661: }
663: th->status = TS_STEP_COMPLETE;
664: PetscFunctionReturn(PETSC_SUCCESS);
665: }
667: static PetscErrorCode TSInterpolate_Theta(TS ts, PetscReal t, Vec X)
668: {
669: TS_Theta *th = (TS_Theta *)ts->data;
670: PetscReal dt = t - ts->ptime;
672: PetscFunctionBegin;
673: PetscCall(VecCopy(ts->vec_sol, th->X));
674: if (th->endpoint) dt *= th->Theta;
675: PetscCall(VecWAXPY(X, dt, th->Xdot, th->X));
676: PetscFunctionReturn(PETSC_SUCCESS);
677: }
679: static PetscErrorCode TSEvaluateWLTE_Theta(TS ts, NormType wnormtype, PetscInt *order, PetscReal *wlte)
680: {
681: TS_Theta *th = (TS_Theta *)ts->data;
682: Vec X = ts->vec_sol; /* X = solution */
683: Vec Y = th->vec_lte_work; /* Y = X + LTE */
684: PetscReal wltea, wlter;
686: PetscFunctionBegin;
687: if (!th->vec_sol_prev) {
688: *wlte = -1;
689: PetscFunctionReturn(PETSC_SUCCESS);
690: }
691: /* Cannot compute LTE in first step or in restart after event */
692: if (ts->steprestart) {
693: *wlte = -1;
694: PetscFunctionReturn(PETSC_SUCCESS);
695: }
696: /* Compute LTE using backward differences with non-constant time step */
697: {
698: PetscReal h = ts->time_step, h_prev = ts->ptime - ts->ptime_prev;
699: PetscReal a = 1 + h_prev / h;
700: PetscScalar scal[3];
701: Vec vecs[3];
702: scal[0] = +1 / a;
703: scal[1] = -1 / (a - 1);
704: scal[2] = +1 / (a * (a - 1));
705: vecs[0] = X;
706: vecs[1] = th->X0;
707: vecs[2] = th->vec_sol_prev;
708: PetscCall(VecCopy(X, Y));
709: PetscCall(VecMAXPY(Y, 3, scal, vecs));
710: PetscCall(TSErrorWeightedNorm(ts, X, Y, wnormtype, wlte, &wltea, &wlter));
711: }
712: if (order) *order = 2;
713: PetscFunctionReturn(PETSC_SUCCESS);
714: }
716: static PetscErrorCode TSRollBack_Theta(TS ts)
717: {
718: TS_Theta *th = (TS_Theta *)ts->data;
719: TS quadts = ts->quadraturets;
721: PetscFunctionBegin;
722: PetscCall(VecCopy(th->X0, ts->vec_sol));
723: if (quadts && ts->costintegralfwd) PetscCall(VecCopy(th->VecCostIntegral0, quadts->vec_sol));
724: th->status = TS_STEP_INCOMPLETE;
725: if (ts->mat_sensip) PetscCall(MatCopy(th->MatFwdSensip0, ts->mat_sensip, SAME_NONZERO_PATTERN));
726: if (quadts && quadts->mat_sensip) PetscCall(MatCopy(th->MatIntegralSensip0, quadts->mat_sensip, SAME_NONZERO_PATTERN));
727: PetscFunctionReturn(PETSC_SUCCESS);
728: }
730: static PetscErrorCode TSForwardStep_Theta(TS ts)
731: {
732: TS_Theta *th = (TS_Theta *)ts->data;
733: TS quadts = ts->quadraturets;
734: Mat MatDeltaFwdSensip = th->MatDeltaFwdSensip;
735: Vec VecDeltaFwdSensipCol = th->VecDeltaFwdSensipCol;
736: PetscInt ntlm;
737: KSP ksp;
738: Mat J, Jpre, quadJ = NULL, quadJp = NULL;
739: PetscScalar *barr, *xarr;
740: PetscReal previous_shift;
742: PetscFunctionBegin;
743: previous_shift = th->shift;
744: PetscCall(MatCopy(ts->mat_sensip, th->MatFwdSensip0, SAME_NONZERO_PATTERN));
746: if (quadts && quadts->mat_sensip) PetscCall(MatCopy(quadts->mat_sensip, th->MatIntegralSensip0, SAME_NONZERO_PATTERN));
747: PetscCall(SNESGetKSP(ts->snes, &ksp));
748: PetscCall(TSGetIJacobian(ts, &J, &Jpre, NULL, NULL));
749: if (quadts) {
750: PetscCall(TSGetRHSJacobian(quadts, &quadJ, NULL, NULL, NULL));
751: PetscCall(TSGetRHSJacobianP(quadts, &quadJp, NULL, NULL));
752: }
754: /* Build RHS */
755: if (th->endpoint) { /* 2-stage method*/
756: th->shift = 1. / ((th->Theta - 1.) * th->time_step0);
757: PetscCall(TSComputeIJacobian(ts, th->ptime0, th->X0, th->Xdot, th->shift, J, Jpre, PETSC_FALSE));
758: PetscCall(MatMatMult(J, ts->mat_sensip, MAT_REUSE_MATRIX, PETSC_DEFAULT, &MatDeltaFwdSensip));
759: PetscCall(MatScale(MatDeltaFwdSensip, (th->Theta - 1.) / th->Theta));
761: /* Add the f_p forcing terms */
762: if (ts->Jacp) {
763: PetscCall(VecZeroEntries(th->Xdot));
764: PetscCall(TSComputeIJacobianP(ts, th->ptime0, th->X0, th->Xdot, th->shift, ts->Jacp, PETSC_FALSE));
765: PetscCall(MatAXPY(MatDeltaFwdSensip, (th->Theta - 1.) / th->Theta, ts->Jacp, SUBSET_NONZERO_PATTERN));
766: th->shift = previous_shift;
767: PetscCall(VecAXPBYPCZ(th->Xdot, -th->shift, th->shift, 0, th->X0, ts->vec_sol));
768: PetscCall(TSComputeIJacobianP(ts, th->stage_time, ts->vec_sol, th->Xdot, th->shift, ts->Jacp, PETSC_FALSE));
769: PetscCall(MatAXPY(MatDeltaFwdSensip, -1., ts->Jacp, SUBSET_NONZERO_PATTERN));
770: }
771: } else { /* 1-stage method */
772: th->shift = 0.0;
773: PetscCall(TSComputeIJacobian(ts, th->stage_time, th->X, th->Xdot, th->shift, J, Jpre, PETSC_FALSE));
774: PetscCall(MatMatMult(J, ts->mat_sensip, MAT_REUSE_MATRIX, PETSC_DEFAULT, &MatDeltaFwdSensip));
775: PetscCall(MatScale(MatDeltaFwdSensip, -1.));
777: /* Add the f_p forcing terms */
778: if (ts->Jacp) {
779: th->shift = previous_shift;
780: PetscCall(VecAXPBYPCZ(th->Xdot, -th->shift, th->shift, 0, th->X0, th->X));
781: PetscCall(TSComputeIJacobianP(ts, th->stage_time, th->X, th->Xdot, th->shift, ts->Jacp, PETSC_FALSE));
782: PetscCall(MatAXPY(MatDeltaFwdSensip, -1., ts->Jacp, SUBSET_NONZERO_PATTERN));
783: }
784: }
786: /* Build LHS */
787: th->shift = previous_shift; /* recover the previous shift used in TSStep_Theta() */
788: if (th->endpoint) {
789: PetscCall(TSComputeIJacobian(ts, th->stage_time, ts->vec_sol, th->Xdot, th->shift, J, Jpre, PETSC_FALSE));
790: } else {
791: PetscCall(TSComputeIJacobian(ts, th->stage_time, th->X, th->Xdot, th->shift, J, Jpre, PETSC_FALSE));
792: }
793: PetscCall(KSPSetOperators(ksp, J, Jpre));
795: /*
796: Evaluate the first stage of integral gradients with the 2-stage method:
797: drdu|t_n*S(t_n) + drdp|t_n
798: This is done before the linear solve because the sensitivity variable S(t_n) will be propagated to S(t_{n+1})
799: */
800: if (th->endpoint) { /* 2-stage method only */
801: if (quadts && quadts->mat_sensip) {
802: PetscCall(TSComputeRHSJacobian(quadts, th->ptime0, th->X0, quadJ, NULL));
803: PetscCall(TSComputeRHSJacobianP(quadts, th->ptime0, th->X0, quadJp));
804: PetscCall(MatTransposeMatMult(ts->mat_sensip, quadJ, MAT_REUSE_MATRIX, PETSC_DEFAULT, &th->MatIntegralSensipTemp));
805: PetscCall(MatAXPY(th->MatIntegralSensipTemp, 1, quadJp, SAME_NONZERO_PATTERN));
806: PetscCall(MatAXPY(quadts->mat_sensip, th->time_step0 * (1. - th->Theta), th->MatIntegralSensipTemp, SAME_NONZERO_PATTERN));
807: }
808: }
810: /* Solve the tangent linear equation for forward sensitivities to parameters */
811: for (ntlm = 0; ntlm < th->num_tlm; ntlm++) {
812: KSPConvergedReason kspreason;
813: PetscCall(MatDenseGetColumn(MatDeltaFwdSensip, ntlm, &barr));
814: PetscCall(VecPlaceArray(VecDeltaFwdSensipCol, barr));
815: if (th->endpoint) {
816: PetscCall(MatDenseGetColumn(ts->mat_sensip, ntlm, &xarr));
817: PetscCall(VecPlaceArray(ts->vec_sensip_col, xarr));
818: PetscCall(KSPSolve(ksp, VecDeltaFwdSensipCol, ts->vec_sensip_col));
819: PetscCall(VecResetArray(ts->vec_sensip_col));
820: PetscCall(MatDenseRestoreColumn(ts->mat_sensip, &xarr));
821: } else {
822: PetscCall(KSPSolve(ksp, VecDeltaFwdSensipCol, VecDeltaFwdSensipCol));
823: }
824: PetscCall(KSPGetConvergedReason(ksp, &kspreason));
825: if (kspreason < 0) {
826: ts->reason = TSFORWARD_DIVERGED_LINEAR_SOLVE;
827: PetscCall(PetscInfo(ts, "Step=%" PetscInt_FMT ", %" PetscInt_FMT "th tangent linear solve, linear solve fails, stopping tangent linear solve\n", ts->steps, ntlm));
828: }
829: PetscCall(VecResetArray(VecDeltaFwdSensipCol));
830: PetscCall(MatDenseRestoreColumn(MatDeltaFwdSensip, &barr));
831: }
833: /*
834: Evaluate the second stage of integral gradients with the 2-stage method:
835: drdu|t_{n+1}*S(t_{n+1}) + drdp|t_{n+1}
836: */
837: if (quadts && quadts->mat_sensip) {
838: if (!th->endpoint) {
839: PetscCall(MatAXPY(ts->mat_sensip, 1, MatDeltaFwdSensip, SAME_NONZERO_PATTERN)); /* stage sensitivity */
840: PetscCall(TSComputeRHSJacobian(quadts, th->stage_time, th->X, quadJ, NULL));
841: PetscCall(TSComputeRHSJacobianP(quadts, th->stage_time, th->X, quadJp));
842: PetscCall(MatTransposeMatMult(ts->mat_sensip, quadJ, MAT_REUSE_MATRIX, PETSC_DEFAULT, &th->MatIntegralSensipTemp));
843: PetscCall(MatAXPY(th->MatIntegralSensipTemp, 1, quadJp, SAME_NONZERO_PATTERN));
844: PetscCall(MatAXPY(quadts->mat_sensip, th->time_step0, th->MatIntegralSensipTemp, SAME_NONZERO_PATTERN));
845: PetscCall(MatAXPY(ts->mat_sensip, (1. - th->Theta) / th->Theta, MatDeltaFwdSensip, SAME_NONZERO_PATTERN));
846: } else {
847: PetscCall(TSComputeRHSJacobian(quadts, th->stage_time, ts->vec_sol, quadJ, NULL));
848: PetscCall(TSComputeRHSJacobianP(quadts, th->stage_time, ts->vec_sol, quadJp));
849: PetscCall(MatTransposeMatMult(ts->mat_sensip, quadJ, MAT_REUSE_MATRIX, PETSC_DEFAULT, &th->MatIntegralSensipTemp));
850: PetscCall(MatAXPY(th->MatIntegralSensipTemp, 1, quadJp, SAME_NONZERO_PATTERN));
851: PetscCall(MatAXPY(quadts->mat_sensip, th->time_step0 * th->Theta, th->MatIntegralSensipTemp, SAME_NONZERO_PATTERN));
852: }
853: } else {
854: if (!th->endpoint) PetscCall(MatAXPY(ts->mat_sensip, 1. / th->Theta, MatDeltaFwdSensip, SAME_NONZERO_PATTERN));
855: }
856: PetscFunctionReturn(PETSC_SUCCESS);
857: }
859: static PetscErrorCode TSForwardGetStages_Theta(TS ts, PetscInt *ns, Mat *stagesensip[])
860: {
861: TS_Theta *th = (TS_Theta *)ts->data;
863: PetscFunctionBegin;
864: if (ns) {
865: if (!th->endpoint && th->Theta != 1.0) *ns = 1; /* midpoint form */
866: else *ns = 2; /* endpoint form */
867: }
868: if (stagesensip) {
869: if (!th->endpoint && th->Theta != 1.0) {
870: th->MatFwdStages[0] = th->MatDeltaFwdSensip;
871: } else {
872: th->MatFwdStages[0] = th->MatFwdSensip0;
873: th->MatFwdStages[1] = ts->mat_sensip; /* stiffly accurate */
874: }
875: *stagesensip = th->MatFwdStages;
876: }
877: PetscFunctionReturn(PETSC_SUCCESS);
878: }
880: /*------------------------------------------------------------*/
881: static PetscErrorCode TSReset_Theta(TS ts)
882: {
883: TS_Theta *th = (TS_Theta *)ts->data;
885: PetscFunctionBegin;
886: PetscCall(VecDestroy(&th->X));
887: PetscCall(VecDestroy(&th->Xdot));
888: PetscCall(VecDestroy(&th->X0));
889: PetscCall(VecDestroy(&th->affine));
891: PetscCall(VecDestroy(&th->vec_sol_prev));
892: PetscCall(VecDestroy(&th->vec_lte_work));
894: PetscCall(VecDestroy(&th->VecCostIntegral0));
895: PetscFunctionReturn(PETSC_SUCCESS);
896: }
898: static PetscErrorCode TSAdjointReset_Theta(TS ts)
899: {
900: TS_Theta *th = (TS_Theta *)ts->data;
902: PetscFunctionBegin;
903: PetscCall(VecDestroyVecs(ts->numcost, &th->VecsDeltaLam));
904: PetscCall(VecDestroyVecs(ts->numcost, &th->VecsDeltaMu));
905: PetscCall(VecDestroyVecs(ts->numcost, &th->VecsDeltaLam2));
906: PetscCall(VecDestroyVecs(ts->numcost, &th->VecsDeltaMu2));
907: PetscCall(VecDestroyVecs(ts->numcost, &th->VecsSensiTemp));
908: PetscCall(VecDestroyVecs(ts->numcost, &th->VecsSensi2Temp));
909: PetscFunctionReturn(PETSC_SUCCESS);
910: }
912: static PetscErrorCode TSDestroy_Theta(TS ts)
913: {
914: PetscFunctionBegin;
915: PetscCall(TSReset_Theta(ts));
916: if (ts->dm) {
917: PetscCall(DMCoarsenHookRemove(ts->dm, DMCoarsenHook_TSTheta, DMRestrictHook_TSTheta, ts));
918: PetscCall(DMSubDomainHookRemove(ts->dm, DMSubDomainHook_TSTheta, DMSubDomainRestrictHook_TSTheta, ts));
919: }
920: PetscCall(PetscFree(ts->data));
921: PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaGetTheta_C", NULL));
922: PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaSetTheta_C", NULL));
923: PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaGetEndpoint_C", NULL));
924: PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaSetEndpoint_C", NULL));
925: PetscFunctionReturn(PETSC_SUCCESS);
926: }
928: /*
929: This defines the nonlinear equation that is to be solved with SNES
930: G(U) = F[t0+Theta*dt, U, (U-U0)*shift] = 0
932: Note that U here is the stage argument. This means that U = U_{n+1} only if endpoint = true,
933: otherwise U = theta U_{n+1} + (1 - theta) U0, which for the case of implicit midpoint is
934: U = (U_{n+1} + U0)/2
935: */
936: static PetscErrorCode SNESTSFormFunction_Theta(SNES snes, Vec x, Vec y, TS ts)
937: {
938: TS_Theta *th = (TS_Theta *)ts->data;
939: Vec X0, Xdot;
940: DM dm, dmsave;
941: PetscReal shift = th->shift;
943: PetscFunctionBegin;
944: PetscCall(SNESGetDM(snes, &dm));
945: /* When using the endpoint variant, this is actually 1/Theta * Xdot */
946: PetscCall(TSThetaGetX0AndXdot(ts, dm, &X0, &Xdot));
947: if (x != X0) {
948: PetscCall(VecAXPBYPCZ(Xdot, -shift, shift, 0, X0, x));
949: } else {
950: PetscCall(VecZeroEntries(Xdot));
951: }
952: /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */
953: dmsave = ts->dm;
954: ts->dm = dm;
955: PetscCall(TSComputeIFunction(ts, th->stage_time, x, Xdot, y, PETSC_FALSE));
956: ts->dm = dmsave;
957: PetscCall(TSThetaRestoreX0AndXdot(ts, dm, &X0, &Xdot));
958: PetscFunctionReturn(PETSC_SUCCESS);
959: }
961: static PetscErrorCode SNESTSFormJacobian_Theta(SNES snes, Vec x, Mat A, Mat B, TS ts)
962: {
963: TS_Theta *th = (TS_Theta *)ts->data;
964: Vec Xdot;
965: DM dm, dmsave;
966: PetscReal shift = th->shift;
968: PetscFunctionBegin;
969: PetscCall(SNESGetDM(snes, &dm));
970: /* Xdot has already been computed in SNESTSFormFunction_Theta (SNES guarantees this) */
971: PetscCall(TSThetaGetX0AndXdot(ts, dm, NULL, &Xdot));
973: dmsave = ts->dm;
974: ts->dm = dm;
975: PetscCall(TSComputeIJacobian(ts, th->stage_time, x, Xdot, shift, A, B, PETSC_FALSE));
976: ts->dm = dmsave;
977: PetscCall(TSThetaRestoreX0AndXdot(ts, dm, NULL, &Xdot));
978: PetscFunctionReturn(PETSC_SUCCESS);
979: }
981: static PetscErrorCode TSForwardSetUp_Theta(TS ts)
982: {
983: TS_Theta *th = (TS_Theta *)ts->data;
984: TS quadts = ts->quadraturets;
986: PetscFunctionBegin;
987: /* combine sensitivities to parameters and sensitivities to initial values into one array */
988: th->num_tlm = ts->num_parameters;
989: PetscCall(MatDuplicate(ts->mat_sensip, MAT_DO_NOT_COPY_VALUES, &th->MatDeltaFwdSensip));
990: if (quadts && quadts->mat_sensip) {
991: PetscCall(MatDuplicate(quadts->mat_sensip, MAT_DO_NOT_COPY_VALUES, &th->MatIntegralSensipTemp));
992: PetscCall(MatDuplicate(quadts->mat_sensip, MAT_DO_NOT_COPY_VALUES, &th->MatIntegralSensip0));
993: }
994: /* backup sensitivity results for roll-backs */
995: PetscCall(MatDuplicate(ts->mat_sensip, MAT_DO_NOT_COPY_VALUES, &th->MatFwdSensip0));
997: PetscCall(VecDuplicate(ts->vec_sol, &th->VecDeltaFwdSensipCol));
998: PetscFunctionReturn(PETSC_SUCCESS);
999: }
1001: static PetscErrorCode TSForwardReset_Theta(TS ts)
1002: {
1003: TS_Theta *th = (TS_Theta *)ts->data;
1004: TS quadts = ts->quadraturets;
1006: PetscFunctionBegin;
1007: if (quadts && quadts->mat_sensip) {
1008: PetscCall(MatDestroy(&th->MatIntegralSensipTemp));
1009: PetscCall(MatDestroy(&th->MatIntegralSensip0));
1010: }
1011: PetscCall(VecDestroy(&th->VecDeltaFwdSensipCol));
1012: PetscCall(MatDestroy(&th->MatDeltaFwdSensip));
1013: PetscCall(MatDestroy(&th->MatFwdSensip0));
1014: PetscFunctionReturn(PETSC_SUCCESS);
1015: }
1017: static PetscErrorCode TSSetUp_Theta(TS ts)
1018: {
1019: TS_Theta *th = (TS_Theta *)ts->data;
1020: TS quadts = ts->quadraturets;
1021: PetscBool match;
1023: PetscFunctionBegin;
1024: if (!th->VecCostIntegral0 && quadts && ts->costintegralfwd) { /* back up cost integral */
1025: PetscCall(VecDuplicate(quadts->vec_sol, &th->VecCostIntegral0));
1026: }
1027: if (!th->X) PetscCall(VecDuplicate(ts->vec_sol, &th->X));
1028: if (!th->Xdot) PetscCall(VecDuplicate(ts->vec_sol, &th->Xdot));
1029: if (!th->X0) PetscCall(VecDuplicate(ts->vec_sol, &th->X0));
1030: if (th->endpoint) PetscCall(VecDuplicate(ts->vec_sol, &th->affine));
1032: th->order = (th->Theta == 0.5) ? 2 : 1;
1033: th->shift = 1 / (th->Theta * ts->time_step);
1035: PetscCall(TSGetDM(ts, &ts->dm));
1036: PetscCall(DMCoarsenHookAdd(ts->dm, DMCoarsenHook_TSTheta, DMRestrictHook_TSTheta, ts));
1037: PetscCall(DMSubDomainHookAdd(ts->dm, DMSubDomainHook_TSTheta, DMSubDomainRestrictHook_TSTheta, ts));
1039: PetscCall(TSGetAdapt(ts, &ts->adapt));
1040: PetscCall(TSAdaptCandidatesClear(ts->adapt));
1041: PetscCall(PetscObjectTypeCompare((PetscObject)ts->adapt, TSADAPTNONE, &match));
1042: if (!match) {
1043: PetscCall(VecDuplicate(ts->vec_sol, &th->vec_sol_prev));
1044: PetscCall(VecDuplicate(ts->vec_sol, &th->vec_lte_work));
1045: }
1046: PetscCall(TSGetSNES(ts, &ts->snes));
1048: ts->stifflyaccurate = (!th->endpoint && th->Theta != 1.0) ? PETSC_FALSE : PETSC_TRUE;
1049: PetscFunctionReturn(PETSC_SUCCESS);
1050: }
1052: /*------------------------------------------------------------*/
1054: static PetscErrorCode TSAdjointSetUp_Theta(TS ts)
1055: {
1056: TS_Theta *th = (TS_Theta *)ts->data;
1058: PetscFunctionBegin;
1059: PetscCall(VecDuplicateVecs(ts->vecs_sensi[0], ts->numcost, &th->VecsDeltaLam));
1060: PetscCall(VecDuplicateVecs(ts->vecs_sensi[0], ts->numcost, &th->VecsSensiTemp));
1061: if (ts->vecs_sensip) PetscCall(VecDuplicateVecs(ts->vecs_sensip[0], ts->numcost, &th->VecsDeltaMu));
1062: if (ts->vecs_sensi2) {
1063: PetscCall(VecDuplicateVecs(ts->vecs_sensi[0], ts->numcost, &th->VecsDeltaLam2));
1064: PetscCall(VecDuplicateVecs(ts->vecs_sensi2[0], ts->numcost, &th->VecsSensi2Temp));
1065: /* hack ts to make implicit TS solver work when users provide only explicit versions of callbacks (RHSFunction,RHSJacobian,RHSHessian etc.) */
1066: if (!ts->ihessianproduct_fuu) ts->vecs_fuu = ts->vecs_guu;
1067: if (!ts->ihessianproduct_fup) ts->vecs_fup = ts->vecs_gup;
1068: }
1069: if (ts->vecs_sensi2p) {
1070: PetscCall(VecDuplicateVecs(ts->vecs_sensi2p[0], ts->numcost, &th->VecsDeltaMu2));
1071: /* hack ts to make implicit TS solver work when users provide only explicit versions of callbacks (RHSFunction,RHSJacobian,RHSHessian etc.) */
1072: if (!ts->ihessianproduct_fpu) ts->vecs_fpu = ts->vecs_gpu;
1073: if (!ts->ihessianproduct_fpp) ts->vecs_fpp = ts->vecs_gpp;
1074: }
1075: PetscFunctionReturn(PETSC_SUCCESS);
1076: }
1078: static PetscErrorCode TSSetFromOptions_Theta(TS ts, PetscOptionItems *PetscOptionsObject)
1079: {
1080: TS_Theta *th = (TS_Theta *)ts->data;
1082: PetscFunctionBegin;
1083: PetscOptionsHeadBegin(PetscOptionsObject, "Theta ODE solver options");
1084: {
1085: PetscCall(PetscOptionsReal("-ts_theta_theta", "Location of stage (0<Theta<=1)", "TSThetaSetTheta", th->Theta, &th->Theta, NULL));
1086: PetscCall(PetscOptionsBool("-ts_theta_endpoint", "Use the endpoint instead of midpoint form of the Theta method", "TSThetaSetEndpoint", th->endpoint, &th->endpoint, NULL));
1087: PetscCall(PetscOptionsBool("-ts_theta_initial_guess_extrapolate", "Extrapolate stage initial guess from previous solution (sometimes unstable)", "TSThetaSetExtrapolate", th->extrapolate, &th->extrapolate, NULL));
1088: }
1089: PetscOptionsHeadEnd();
1090: PetscFunctionReturn(PETSC_SUCCESS);
1091: }
1093: static PetscErrorCode TSView_Theta(TS ts, PetscViewer viewer)
1094: {
1095: TS_Theta *th = (TS_Theta *)ts->data;
1096: PetscBool iascii;
1098: PetscFunctionBegin;
1099: PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &iascii));
1100: if (iascii) {
1101: PetscCall(PetscViewerASCIIPrintf(viewer, " Theta=%g\n", (double)th->Theta));
1102: PetscCall(PetscViewerASCIIPrintf(viewer, " Extrapolation=%s\n", th->extrapolate ? "yes" : "no"));
1103: }
1104: PetscFunctionReturn(PETSC_SUCCESS);
1105: }
1107: static PetscErrorCode TSThetaGetTheta_Theta(TS ts, PetscReal *theta)
1108: {
1109: TS_Theta *th = (TS_Theta *)ts->data;
1111: PetscFunctionBegin;
1112: *theta = th->Theta;
1113: PetscFunctionReturn(PETSC_SUCCESS);
1114: }
1116: static PetscErrorCode TSThetaSetTheta_Theta(TS ts, PetscReal theta)
1117: {
1118: TS_Theta *th = (TS_Theta *)ts->data;
1120: PetscFunctionBegin;
1121: PetscCheck(theta > 0 && theta <= 1, PetscObjectComm((PetscObject)ts), PETSC_ERR_ARG_OUTOFRANGE, "Theta %g not in range (0,1]", (double)theta);
1122: th->Theta = theta;
1123: th->order = (th->Theta == 0.5) ? 2 : 1;
1124: PetscFunctionReturn(PETSC_SUCCESS);
1125: }
1127: static PetscErrorCode TSThetaGetEndpoint_Theta(TS ts, PetscBool *endpoint)
1128: {
1129: TS_Theta *th = (TS_Theta *)ts->data;
1131: PetscFunctionBegin;
1132: *endpoint = th->endpoint;
1133: PetscFunctionReturn(PETSC_SUCCESS);
1134: }
1136: static PetscErrorCode TSThetaSetEndpoint_Theta(TS ts, PetscBool flg)
1137: {
1138: TS_Theta *th = (TS_Theta *)ts->data;
1140: PetscFunctionBegin;
1141: th->endpoint = flg;
1142: PetscFunctionReturn(PETSC_SUCCESS);
1143: }
1145: #if defined(PETSC_HAVE_COMPLEX)
1146: static PetscErrorCode TSComputeLinearStability_Theta(TS ts, PetscReal xr, PetscReal xi, PetscReal *yr, PetscReal *yi)
1147: {
1148: PetscComplex z = xr + xi * PETSC_i, f;
1149: TS_Theta *th = (TS_Theta *)ts->data;
1150: const PetscReal one = 1.0;
1152: PetscFunctionBegin;
1153: f = (one + (one - th->Theta) * z) / (one - th->Theta * z);
1154: *yr = PetscRealPartComplex(f);
1155: *yi = PetscImaginaryPartComplex(f);
1156: PetscFunctionReturn(PETSC_SUCCESS);
1157: }
1158: #endif
1160: static PetscErrorCode TSGetStages_Theta(TS ts, PetscInt *ns, Vec *Y[])
1161: {
1162: TS_Theta *th = (TS_Theta *)ts->data;
1164: PetscFunctionBegin;
1165: if (ns) {
1166: if (!th->endpoint && th->Theta != 1.0) *ns = 1; /* midpoint form */
1167: else *ns = 2; /* endpoint form */
1168: }
1169: if (Y) {
1170: if (!th->endpoint && th->Theta != 1.0) {
1171: th->Stages[0] = th->X;
1172: } else {
1173: th->Stages[0] = th->X0;
1174: th->Stages[1] = ts->vec_sol; /* stiffly accurate */
1175: }
1176: *Y = th->Stages;
1177: }
1178: PetscFunctionReturn(PETSC_SUCCESS);
1179: }
1181: /* ------------------------------------------------------------ */
1182: /*MC
1183: TSTHETA - DAE solver using the implicit Theta method
1185: Level: beginner
1187: Options Database Keys:
1188: + -ts_theta_theta <Theta> - Location of stage (0<Theta<=1)
1189: . -ts_theta_endpoint <flag> - Use the endpoint (like Crank-Nicholson) instead of midpoint form of the Theta method
1190: - -ts_theta_initial_guess_extrapolate <flg> - Extrapolate stage initial guess from previous solution (sometimes unstable)
1192: Notes:
1193: .vb
1194: -ts_type theta -ts_theta_theta 1.0 corresponds to backward Euler (TSBEULER)
1195: -ts_type theta -ts_theta_theta 0.5 corresponds to the implicit midpoint rule
1196: -ts_type theta -ts_theta_theta 0.5 -ts_theta_endpoint corresponds to Crank-Nicholson (TSCN)
1197: .ve
1199: The endpoint variant of the Theta method and backward Euler can be applied to DAE. The midpoint variant is not suitable for DAEs because it is not stiffly accurate.
1201: The midpoint variant is cast as a 1-stage implicit Runge-Kutta method.
1203: .vb
1204: Theta | Theta
1205: -------------
1206: | 1
1207: .ve
1209: For the default Theta=0.5, this is also known as the implicit midpoint rule.
1211: When the endpoint variant is chosen, the method becomes a 2-stage method with first stage explicit:
1213: .vb
1214: 0 | 0 0
1215: 1 | 1-Theta Theta
1216: -------------------
1217: | 1-Theta Theta
1218: .ve
1220: For the default Theta=0.5, this is the trapezoid rule (also known as Crank-Nicolson, see TSCN).
1222: To apply a diagonally implicit RK method to DAE, the stage formula
1224: $ Y_i = X + h sum_j a_ij Y'_j
1226: is interpreted as a formula for Y'_i in terms of Y_i and known values (Y'_j, j<i)
1228: .seealso: [](ch_ts), `TSCreate()`, `TS`, `TSSetType()`, `TSCN`, `TSBEULER`, `TSThetaSetTheta()`, `TSThetaSetEndpoint()`
1229: M*/
1230: PETSC_EXTERN PetscErrorCode TSCreate_Theta(TS ts)
1231: {
1232: TS_Theta *th;
1234: PetscFunctionBegin;
1235: ts->ops->reset = TSReset_Theta;
1236: ts->ops->adjointreset = TSAdjointReset_Theta;
1237: ts->ops->destroy = TSDestroy_Theta;
1238: ts->ops->view = TSView_Theta;
1239: ts->ops->setup = TSSetUp_Theta;
1240: ts->ops->adjointsetup = TSAdjointSetUp_Theta;
1241: ts->ops->adjointreset = TSAdjointReset_Theta;
1242: ts->ops->step = TSStep_Theta;
1243: ts->ops->interpolate = TSInterpolate_Theta;
1244: ts->ops->evaluatewlte = TSEvaluateWLTE_Theta;
1245: ts->ops->rollback = TSRollBack_Theta;
1246: ts->ops->resizeregister = TSResizeRegister_Theta;
1247: ts->ops->setfromoptions = TSSetFromOptions_Theta;
1248: ts->ops->snesfunction = SNESTSFormFunction_Theta;
1249: ts->ops->snesjacobian = SNESTSFormJacobian_Theta;
1250: #if defined(PETSC_HAVE_COMPLEX)
1251: ts->ops->linearstability = TSComputeLinearStability_Theta;
1252: #endif
1253: ts->ops->getstages = TSGetStages_Theta;
1254: ts->ops->adjointstep = TSAdjointStep_Theta;
1255: ts->ops->adjointintegral = TSAdjointCostIntegral_Theta;
1256: ts->ops->forwardintegral = TSForwardCostIntegral_Theta;
1257: ts->default_adapt_type = TSADAPTNONE;
1259: ts->ops->forwardsetup = TSForwardSetUp_Theta;
1260: ts->ops->forwardreset = TSForwardReset_Theta;
1261: ts->ops->forwardstep = TSForwardStep_Theta;
1262: ts->ops->forwardgetstages = TSForwardGetStages_Theta;
1264: ts->usessnes = PETSC_TRUE;
1266: PetscCall(PetscNew(&th));
1267: ts->data = (void *)th;
1269: th->VecsDeltaLam = NULL;
1270: th->VecsDeltaMu = NULL;
1271: th->VecsSensiTemp = NULL;
1272: th->VecsSensi2Temp = NULL;
1274: th->extrapolate = PETSC_FALSE;
1275: th->Theta = 0.5;
1276: th->order = 2;
1277: PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaGetTheta_C", TSThetaGetTheta_Theta));
1278: PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaSetTheta_C", TSThetaSetTheta_Theta));
1279: PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaGetEndpoint_C", TSThetaGetEndpoint_Theta));
1280: PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaSetEndpoint_C", TSThetaSetEndpoint_Theta));
1281: PetscFunctionReturn(PETSC_SUCCESS);
1282: }
1284: /*@
1285: TSThetaGetTheta - Get the abscissa of the stage in (0,1] for `TSTHETA`
1287: Not Collective
1289: Input Parameter:
1290: . ts - timestepping context
1292: Output Parameter:
1293: . theta - stage abscissa
1295: Level: advanced
1297: Note:
1298: Use of this function is normally only required to hack `TSTHETA` to use a modified integration scheme.
1300: .seealso: [](ch_ts), `TSThetaSetTheta()`, `TSTHETA`
1301: @*/
1302: PetscErrorCode TSThetaGetTheta(TS ts, PetscReal *theta)
1303: {
1304: PetscFunctionBegin;
1306: PetscAssertPointer(theta, 2);
1307: PetscUseMethod(ts, "TSThetaGetTheta_C", (TS, PetscReal *), (ts, theta));
1308: PetscFunctionReturn(PETSC_SUCCESS);
1309: }
1311: /*@
1312: TSThetaSetTheta - Set the abscissa of the stage in (0,1] for `TSTHETA`
1314: Not Collective
1316: Input Parameters:
1317: + ts - timestepping context
1318: - theta - stage abscissa
1320: Options Database Key:
1321: . -ts_theta_theta <theta> - set theta
1323: Level: intermediate
1325: .seealso: [](ch_ts), `TSThetaGetTheta()`, `TSTHETA`, `TSCN`
1326: @*/
1327: PetscErrorCode TSThetaSetTheta(TS ts, PetscReal theta)
1328: {
1329: PetscFunctionBegin;
1331: PetscTryMethod(ts, "TSThetaSetTheta_C", (TS, PetscReal), (ts, theta));
1332: PetscFunctionReturn(PETSC_SUCCESS);
1333: }
1335: /*@
1336: TSThetaGetEndpoint - Gets whether to use the endpoint variant of the method (e.g. trapezoid/Crank-Nicolson instead of midpoint rule) for `TSTHETA`
1338: Not Collective
1340: Input Parameter:
1341: . ts - timestepping context
1343: Output Parameter:
1344: . endpoint - `PETSC_TRUE` when using the endpoint variant
1346: Level: advanced
1348: .seealso: [](ch_ts), `TSThetaSetEndpoint()`, `TSTHETA`, `TSCN`
1349: @*/
1350: PetscErrorCode TSThetaGetEndpoint(TS ts, PetscBool *endpoint)
1351: {
1352: PetscFunctionBegin;
1354: PetscAssertPointer(endpoint, 2);
1355: PetscUseMethod(ts, "TSThetaGetEndpoint_C", (TS, PetscBool *), (ts, endpoint));
1356: PetscFunctionReturn(PETSC_SUCCESS);
1357: }
1359: /*@
1360: TSThetaSetEndpoint - Sets whether to use the endpoint variant of the method (e.g. trapezoid/Crank-Nicolson instead of midpoint rule) for `TSTHETA`
1362: Not Collective
1364: Input Parameters:
1365: + ts - timestepping context
1366: - flg - `PETSC_TRUE` to use the endpoint variant
1368: Options Database Key:
1369: . -ts_theta_endpoint <flg> - use the endpoint variant
1371: Level: intermediate
1373: .seealso: [](ch_ts), `TSTHETA`, `TSCN`
1374: @*/
1375: PetscErrorCode TSThetaSetEndpoint(TS ts, PetscBool flg)
1376: {
1377: PetscFunctionBegin;
1379: PetscTryMethod(ts, "TSThetaSetEndpoint_C", (TS, PetscBool), (ts, flg));
1380: PetscFunctionReturn(PETSC_SUCCESS);
1381: }
1383: /*
1384: * TSBEULER and TSCN are straightforward specializations of TSTHETA.
1385: * The creation functions for these specializations are below.
1386: */
1388: static PetscErrorCode TSSetUp_BEuler(TS ts)
1389: {
1390: TS_Theta *th = (TS_Theta *)ts->data;
1392: PetscFunctionBegin;
1393: PetscCheck(th->Theta == 1.0, PetscObjectComm((PetscObject)ts), PETSC_ERR_OPT_OVERWRITE, "Can not change the default value (1) of theta when using backward Euler");
1394: PetscCheck(!th->endpoint, PetscObjectComm((PetscObject)ts), PETSC_ERR_OPT_OVERWRITE, "Can not change to the endpoint form of the Theta methods when using backward Euler");
1395: PetscCall(TSSetUp_Theta(ts));
1396: PetscFunctionReturn(PETSC_SUCCESS);
1397: }
1399: static PetscErrorCode TSView_BEuler(TS ts, PetscViewer viewer)
1400: {
1401: PetscFunctionBegin;
1402: PetscFunctionReturn(PETSC_SUCCESS);
1403: }
1405: /*MC
1406: TSBEULER - ODE solver using the implicit backward Euler method
1408: Level: beginner
1410: Note:
1411: `TSBEULER` is equivalent to `TSTHETA` with Theta=1.0 or `-ts_type theta -ts_theta_theta 1.0`
1413: .seealso: [](ch_ts), `TSCreate()`, `TS`, `TSSetType()`, `TSEULER`, `TSCN`, `TSTHETA`
1414: M*/
1415: PETSC_EXTERN PetscErrorCode TSCreate_BEuler(TS ts)
1416: {
1417: PetscFunctionBegin;
1418: PetscCall(TSCreate_Theta(ts));
1419: PetscCall(TSThetaSetTheta(ts, 1.0));
1420: PetscCall(TSThetaSetEndpoint(ts, PETSC_FALSE));
1421: ts->ops->setup = TSSetUp_BEuler;
1422: ts->ops->view = TSView_BEuler;
1423: PetscFunctionReturn(PETSC_SUCCESS);
1424: }
1426: static PetscErrorCode TSSetUp_CN(TS ts)
1427: {
1428: TS_Theta *th = (TS_Theta *)ts->data;
1430: PetscFunctionBegin;
1431: PetscCheck(th->Theta == 0.5, PetscObjectComm((PetscObject)ts), PETSC_ERR_OPT_OVERWRITE, "Can not change the default value (0.5) of theta when using Crank-Nicolson");
1432: PetscCheck(th->endpoint, PetscObjectComm((PetscObject)ts), PETSC_ERR_OPT_OVERWRITE, "Can not change to the midpoint form of the Theta methods when using Crank-Nicolson");
1433: PetscCall(TSSetUp_Theta(ts));
1434: PetscFunctionReturn(PETSC_SUCCESS);
1435: }
1437: static PetscErrorCode TSView_CN(TS ts, PetscViewer viewer)
1438: {
1439: PetscFunctionBegin;
1440: PetscFunctionReturn(PETSC_SUCCESS);
1441: }
1443: /*MC
1444: TSCN - ODE solver using the implicit Crank-Nicolson method.
1446: Level: beginner
1448: Notes:
1449: `TSCN` is equivalent to `TSTHETA` with Theta=0.5 and the "endpoint" option set. I.e.
1450: .vb
1451: -ts_type theta
1452: -ts_theta_theta 0.5
1453: -ts_theta_endpoint
1454: .ve
1456: .seealso: [](ch_ts), `TSCreate()`, `TS`, `TSSetType()`, `TSBEULER`, `TSTHETA`, `TSType`,
1457: M*/
1458: PETSC_EXTERN PetscErrorCode TSCreate_CN(TS ts)
1459: {
1460: PetscFunctionBegin;
1461: PetscCall(TSCreate_Theta(ts));
1462: PetscCall(TSThetaSetTheta(ts, 0.5));
1463: PetscCall(TSThetaSetEndpoint(ts, PETSC_TRUE));
1464: ts->ops->setup = TSSetUp_CN;
1465: ts->ops->view = TSView_CN;
1466: PetscFunctionReturn(PETSC_SUCCESS);
1467: }