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