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