Actual source code: theta.c
petsc-3.12.5 2020-03-29
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 X0,X,Xdot; /* Storage for stages and time derivative */
13: Vec affine; /* Affine vector needed for residual at beginning of step in endpoint formulation */
14: PetscReal Theta;
15: PetscReal ptime;
16: PetscReal time_step;
17: PetscReal shift;
18: PetscInt order;
19: PetscBool endpoint;
20: PetscBool extrapolate;
21: TSStepStatus status;
22: Vec VecCostIntegral0; /* Backup for roll-backs due to events */
24: /* context for sensitivity analysis */
25: PetscInt num_tlm; /* Total number of tangent linear equations */
26: Vec *VecsDeltaLam; /* Increment of the adjoint sensitivity w.r.t IC at stage */
27: Vec *VecsDeltaMu; /* Increment of the adjoint sensitivity w.r.t P at stage */
28: Vec *VecsSensiTemp; /* Vector to be multiplied with Jacobian transpose */
29: Mat MatDeltaFwdSensip; /* Increment of the forward sensitivity at stage */
30: Vec VecDeltaFwdSensipCol; /* Working vector for holding one column of the sensitivity matrix */
31: Mat MatFwdSensip0; /* backup for roll-backs due to events */
32: Mat MatIntegralSensipTemp; /* Working vector for forward integral sensitivity */
33: Mat MatIntegralSensip0; /* backup for roll-backs due to events */
34: Vec *VecsDeltaLam2; /* Increment of the 2nd-order adjoint sensitivity w.r.t IC at stage */
35: Vec *VecsDeltaMu2; /* Increment of the 2nd-order adjoint sensitivity w.r.t P at stage */
36: Vec *VecsSensi2Temp; /* Working vectors that holds the residual for the second-order adjoint */
37: Vec *VecsAffine; /* Working vectors to store residuals */
38: /* context for error estimation */
39: Vec vec_sol_prev;
40: Vec vec_lte_work;
41: } TS_Theta;
43: static PetscErrorCode TSThetaGetX0AndXdot(TS ts,DM dm,Vec *X0,Vec *Xdot)
44: {
45: TS_Theta *th = (TS_Theta*)ts->data;
49: if (X0) {
50: if (dm && dm != ts->dm) {
51: DMGetNamedGlobalVector(dm,"TSTheta_X0",X0);
52: } else *X0 = ts->vec_sol;
53: }
54: if (Xdot) {
55: if (dm && dm != ts->dm) {
56: DMGetNamedGlobalVector(dm,"TSTheta_Xdot",Xdot);
57: } else *Xdot = th->Xdot;
58: }
59: return(0);
60: }
62: static PetscErrorCode TSThetaRestoreX0AndXdot(TS ts,DM dm,Vec *X0,Vec *Xdot)
63: {
67: if (X0) {
68: if (dm && dm != ts->dm) {
69: DMRestoreNamedGlobalVector(dm,"TSTheta_X0",X0);
70: }
71: }
72: if (Xdot) {
73: if (dm && dm != ts->dm) {
74: DMRestoreNamedGlobalVector(dm,"TSTheta_Xdot",Xdot);
75: }
76: }
77: return(0);
78: }
80: static PetscErrorCode DMCoarsenHook_TSTheta(DM fine,DM coarse,void *ctx)
81: {
83: return(0);
84: }
86: static PetscErrorCode DMRestrictHook_TSTheta(DM fine,Mat restrct,Vec rscale,Mat inject,DM coarse,void *ctx)
87: {
88: TS ts = (TS)ctx;
90: Vec X0,Xdot,X0_c,Xdot_c;
93: TSThetaGetX0AndXdot(ts,fine,&X0,&Xdot);
94: TSThetaGetX0AndXdot(ts,coarse,&X0_c,&Xdot_c);
95: MatRestrict(restrct,X0,X0_c);
96: MatRestrict(restrct,Xdot,Xdot_c);
97: VecPointwiseMult(X0_c,rscale,X0_c);
98: VecPointwiseMult(Xdot_c,rscale,Xdot_c);
99: TSThetaRestoreX0AndXdot(ts,fine,&X0,&Xdot);
100: TSThetaRestoreX0AndXdot(ts,coarse,&X0_c,&Xdot_c);
101: return(0);
102: }
104: static PetscErrorCode DMSubDomainHook_TSTheta(DM dm,DM subdm,void *ctx)
105: {
107: return(0);
108: }
110: static PetscErrorCode DMSubDomainRestrictHook_TSTheta(DM dm,VecScatter gscat,VecScatter lscat,DM subdm,void *ctx)
111: {
112: TS ts = (TS)ctx;
114: Vec X0,Xdot,X0_sub,Xdot_sub;
117: TSThetaGetX0AndXdot(ts,dm,&X0,&Xdot);
118: TSThetaGetX0AndXdot(ts,subdm,&X0_sub,&Xdot_sub);
120: VecScatterBegin(gscat,X0,X0_sub,INSERT_VALUES,SCATTER_FORWARD);
121: VecScatterEnd(gscat,X0,X0_sub,INSERT_VALUES,SCATTER_FORWARD);
123: VecScatterBegin(gscat,Xdot,Xdot_sub,INSERT_VALUES,SCATTER_FORWARD);
124: VecScatterEnd(gscat,Xdot,Xdot_sub,INSERT_VALUES,SCATTER_FORWARD);
126: TSThetaRestoreX0AndXdot(ts,dm,&X0,&Xdot);
127: TSThetaRestoreX0AndXdot(ts,subdm,&X0_sub,&Xdot_sub);
128: return(0);
129: }
131: static PetscErrorCode TSThetaEvaluateCostIntegral(TS ts)
132: {
133: TS_Theta *th = (TS_Theta*)ts->data;
134: TS quadts = ts->quadraturets;
138: if (th->endpoint) {
139: /* Evolve ts->vec_costintegral to compute integrals */
140: if (th->Theta!=1.0) {
141: TSComputeRHSFunction(quadts,th->ptime,th->X0,ts->vec_costintegrand);
142: VecAXPY(quadts->vec_sol,th->time_step*(1.0-th->Theta),ts->vec_costintegrand);
143: }
144: TSComputeRHSFunction(quadts,ts->ptime,ts->vec_sol,ts->vec_costintegrand);
145: VecAXPY(quadts->vec_sol,th->time_step*th->Theta,ts->vec_costintegrand);
146: } else {
147: TSComputeRHSFunction(quadts,th->stage_time,th->X,ts->vec_costintegrand);
148: VecAXPY(quadts->vec_sol,th->time_step,ts->vec_costintegrand);
149: }
150: return(0);
151: }
153: static PetscErrorCode TSForwardCostIntegral_Theta(TS ts)
154: {
155: TS_Theta *th = (TS_Theta*)ts->data;
156: TS quadts = ts->quadraturets;
160: /* backup cost integral */
161: VecCopy(quadts->vec_sol,th->VecCostIntegral0);
162: TSThetaEvaluateCostIntegral(ts);
163: return(0);
164: }
166: static PetscErrorCode TSAdjointCostIntegral_Theta(TS ts)
167: {
171: TSThetaEvaluateCostIntegral(ts);
172: return(0);
173: }
175: static PetscErrorCode TSTheta_SNESSolve(TS ts,Vec b,Vec x)
176: {
177: PetscInt nits,lits;
181: SNESSolve(ts->snes,b,x);
182: SNESGetIterationNumber(ts->snes,&nits);
183: SNESGetLinearSolveIterations(ts->snes,&lits);
184: ts->snes_its += nits; ts->ksp_its += lits;
185: return(0);
186: }
188: static PetscErrorCode TSStep_Theta(TS ts)
189: {
190: TS_Theta *th = (TS_Theta*)ts->data;
191: PetscInt rejections = 0;
192: PetscBool stageok,accept = PETSC_TRUE;
193: PetscReal next_time_step = ts->time_step;
197: if (!ts->steprollback) {
198: if (th->vec_sol_prev) { VecCopy(th->X0,th->vec_sol_prev); }
199: VecCopy(ts->vec_sol,th->X0);
200: }
202: th->status = TS_STEP_INCOMPLETE;
203: while (!ts->reason && th->status != TS_STEP_COMPLETE) {
205: th->shift = 1/(th->Theta*ts->time_step);
206: th->stage_time = ts->ptime + (th->endpoint ? (PetscReal)1 : th->Theta)*ts->time_step;
208: VecCopy(th->X0,th->X);
209: if (th->extrapolate && !ts->steprestart) {
210: VecAXPY(th->X,1/th->shift,th->Xdot);
211: }
212: if (th->endpoint) { /* This formulation assumes linear time-independent mass matrix */
213: if (!th->affine) {VecDuplicate(ts->vec_sol,&th->affine);}
214: VecZeroEntries(th->Xdot);
215: TSComputeIFunction(ts,ts->ptime,th->X0,th->Xdot,th->affine,PETSC_FALSE);
216: VecScale(th->affine,(th->Theta-1)/th->Theta);
217: } else if (th->affine) { /* Just in case th->endpoint is changed between calls to TSStep_Theta() */
218: VecZeroEntries(th->affine);
219: }
220: TSPreStage(ts,th->stage_time);
221: TSTheta_SNESSolve(ts,th->affine,th->X);
222: TSPostStage(ts,th->stage_time,0,&th->X);
223: TSAdaptCheckStage(ts->adapt,ts,th->stage_time,th->X,&stageok);
224: if (!stageok) goto reject_step;
226: th->status = TS_STEP_PENDING;
227: if (th->endpoint) {
228: VecCopy(th->X,ts->vec_sol);
229: } else {
230: VecAXPBYPCZ(th->Xdot,-th->shift,th->shift,0,th->X0,th->X);
231: VecAXPY(ts->vec_sol,ts->time_step,th->Xdot);
232: }
233: TSAdaptChoose(ts->adapt,ts,ts->time_step,NULL,&next_time_step,&accept);
234: th->status = accept ? TS_STEP_COMPLETE : TS_STEP_INCOMPLETE;
235: if (!accept) {
236: VecCopy(th->X0,ts->vec_sol);
237: ts->time_step = next_time_step;
238: goto reject_step;
239: }
241: if (ts->forward_solve || ts->costintegralfwd) { /* Save the info for the later use in cost integral evaluation */
242: th->ptime = ts->ptime;
243: th->time_step = ts->time_step;
244: }
246: ts->ptime += ts->time_step;
247: ts->time_step = next_time_step;
248: break;
250: reject_step:
251: ts->reject++; accept = PETSC_FALSE;
252: if (!ts->reason && ++rejections > ts->max_reject && ts->max_reject >= 0) {
253: ts->reason = TS_DIVERGED_STEP_REJECTED;
254: PetscInfo2(ts,"Step=%D, step rejections %D greater than current TS allowed, stopping solve\n",ts->steps,rejections);
255: }
256: }
257: return(0);
258: }
260: static PetscErrorCode TSAdjointStepBEuler_Private(TS ts)
261: {
262: TS_Theta *th = (TS_Theta*)ts->data;
263: TS quadts = ts->quadraturets;
264: Vec *VecsDeltaLam = th->VecsDeltaLam,*VecsDeltaMu = th->VecsDeltaMu,*VecsSensiTemp = th->VecsSensiTemp;
265: Vec *VecsDeltaLam2 = th->VecsDeltaLam2,*VecsDeltaMu2 = th->VecsDeltaMu2,*VecsSensi2Temp = th->VecsSensi2Temp;
266: PetscInt nadj;
267: Mat J,Jpre,quadJ = NULL,quadJp = NULL;
268: KSP ksp;
269: PetscScalar *xarr;
270: TSEquationType eqtype;
271: PetscBool isexplicitode = PETSC_FALSE;
275: TSGetEquationType(ts,&eqtype);
276: if (eqtype == TS_EQ_ODE_EXPLICIT) {
277: isexplicitode = PETSC_TRUE;
278: VecsDeltaLam = ts->vecs_sensi;
279: VecsDeltaLam2 = ts->vecs_sensi2;
280: }
281: th->status = TS_STEP_INCOMPLETE;
282: SNESGetKSP(ts->snes,&ksp);
283: TSGetIJacobian(ts,&J,&Jpre,NULL,NULL);
284: if (quadts) {
285: TSGetRHSJacobian(quadts,&quadJ,NULL,NULL,NULL);
286: TSGetRHSJacobianP(quadts,&quadJp,NULL,NULL);
287: }
289: /* If endpoint=1, th->ptime and th->X0 will be used; if endpoint=0, th->stage_time and th->X will be used. */
290: th->stage_time = ts->ptime; /* time_step is negative*/
291: th->ptime = ts->ptime + ts->time_step;
292: th->time_step = -ts->time_step;
294: /* Build RHS for first-order adjoint lambda_{n+1}/h + r_u^T(n+1) */
295: if (quadts) {
296: TSComputeRHSJacobian(quadts,th->stage_time,th->X,quadJ,NULL);
297: }
299: for (nadj=0; nadj<ts->numcost; nadj++) {
300: VecCopy(ts->vecs_sensi[nadj],VecsSensiTemp[nadj]);
301: VecScale(VecsSensiTemp[nadj],1./th->time_step); /* lambda_{n+1}/h */
302: if (quadJ) {
303: MatDenseGetColumn(quadJ,nadj,&xarr);
304: VecPlaceArray(ts->vec_drdu_col,xarr);
305: VecAXPY(VecsSensiTemp[nadj],1.,ts->vec_drdu_col);
306: VecResetArray(ts->vec_drdu_col);
307: MatDenseRestoreColumn(quadJ,&xarr);
308: }
309: }
311: /* Build LHS for first-order adjoint */
312: TSComputeSNESJacobian(ts,th->X,J,Jpre);
313: KSPSetOperators(ksp,J,Jpre);
315: /* Solve stage equation LHS*lambda_s = RHS for first-order adjoint */
316: for (nadj=0; nadj<ts->numcost; nadj++) {
317: KSPConvergedReason kspreason;
318: KSPSolveTranspose(ksp,VecsSensiTemp[nadj],VecsDeltaLam[nadj]);
319: KSPGetConvergedReason(ksp,&kspreason);
320: if (kspreason < 0) {
321: ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
322: PetscInfo2(ts,"Step=%D, %Dth cost function, transposed linear solve fails, stopping 1st-order adjoint solve\n",ts->steps,nadj);
323: }
324: }
326: if (ts->vecs_sensi2) { /* U_{n+1} */
327: /* Get w1 at t_{n+1} from TLM matrix */
328: MatDenseGetColumn(ts->mat_sensip,0,&xarr);
329: VecPlaceArray(ts->vec_sensip_col,xarr);
330: /* lambda_s^T F_UU w_1 */
331: TSComputeIHessianProductFunctionUU(ts,th->stage_time,th->X,VecsDeltaLam,ts->vec_sensip_col,ts->vecs_fuu);
332: /* lambda_s^T F_UP w_2 */
333: TSComputeIHessianProductFunctionUP(ts,th->stage_time,th->X,VecsDeltaLam,ts->vec_dir,ts->vecs_fup);
334: for (nadj=0; nadj<ts->numcost; nadj++) { /* compute the residual */
335: VecCopy(ts->vecs_sensi2[nadj],VecsSensi2Temp[nadj]);
336: VecScale(VecsSensi2Temp[nadj],1./th->time_step);
337: VecAXPY(VecsSensi2Temp[nadj],-1.,ts->vecs_fuu[nadj]);
338: if (ts->vecs_fup) {
339: VecAXPY(VecsSensi2Temp[nadj],-1.,ts->vecs_fup[nadj]);
340: }
341: }
342: /* Solve stage equation LHS X = RHS for second-order adjoint */
343: for (nadj=0; nadj<ts->numcost; nadj++) {
344: KSPConvergedReason kspreason;
345: KSPSolveTranspose(ksp,VecsSensi2Temp[nadj],VecsDeltaLam2[nadj]);
346: KSPGetConvergedReason(ksp,&kspreason);
347: if (kspreason < 0) {
348: ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
349: PetscInfo2(ts,"Step=%D, %Dth cost function, transposed linear solve fails, stopping 2nd-order adjoint solve\n",ts->steps,nadj);
350: }
351: }
352: }
354: /* Update sensitivities, and evaluate integrals if there is any */
355: if (!isexplicitode) {
356: th->shift = 0.0;
357: TSComputeSNESJacobian(ts,th->X,J,Jpre);
358: KSPSetOperators(ksp,J,Jpre);
359: MatScale(J,-1.);
360: for (nadj=0; nadj<ts->numcost; nadj++) {
361: /* Add f_U \lambda_s to the original RHS */
362: MatMultTransposeAdd(J,VecsDeltaLam[nadj],VecsSensiTemp[nadj],VecsSensiTemp[nadj]);
363: VecScale(VecsSensiTemp[nadj],th->time_step);
364: VecCopy(VecsSensiTemp[nadj],ts->vecs_sensi[nadj]);
365: if (ts->vecs_sensi2) {
366: MatMultTransposeAdd(J,VecsDeltaLam2[nadj],VecsSensi2Temp[nadj],VecsSensi2Temp[nadj]);
367: VecScale(VecsSensi2Temp[nadj],th->time_step);
368: VecCopy(VecsSensi2Temp[nadj],ts->vecs_sensi2[nadj]);
369: }
370: }
371: }
372: if (ts->vecs_sensip) {
373: th->shift = 1./th->time_step;
374: TSComputeIJacobianP(ts,th->stage_time,th->X,th->Xdot,th->shift,ts->Jacp,PETSC_FALSE); /* get -f_p */
375: if (quadts) {
376: TSComputeRHSJacobianP(quadts,th->stage_time,th->X,quadJp);
377: }
378: if (ts->vecs_sensi2p) {
379: /* lambda_s^T F_PU w_1 */
380: TSComputeIHessianProductFunctionPU(ts,th->stage_time,th->X,VecsDeltaLam,ts->vec_sensip_col,ts->vecs_fpu);
381: /* lambda_s^T F_PP w_2 */
382: TSComputeIHessianProductFunctionPP(ts,th->stage_time,th->X,VecsDeltaLam,ts->vec_dir,ts->vecs_fpp);
383: }
385: for (nadj=0; nadj<ts->numcost; nadj++) {
386: MatMultTranspose(ts->Jacp,VecsDeltaLam[nadj],VecsDeltaMu[nadj]);
387: VecAXPY(ts->vecs_sensip[nadj],-th->time_step,VecsDeltaMu[nadj]);
388: if (quadJp) {
389: MatDenseGetColumn(quadJp,nadj,&xarr);
390: VecPlaceArray(ts->vec_drdp_col,xarr);
391: VecAXPY(ts->vecs_sensip[nadj],th->time_step,ts->vec_drdp_col);
392: VecResetArray(ts->vec_drdp_col);
393: MatDenseRestoreColumn(quadJp,&xarr);
394: }
395: if (ts->vecs_sensi2p) {
396: MatMultTranspose(ts->Jacp,VecsDeltaLam2[nadj],VecsDeltaMu2[nadj]);
397: VecAXPY(ts->vecs_sensi2p[nadj],-th->time_step,VecsDeltaMu2[nadj]);
398: if (ts->vecs_fpu) {
399: VecAXPY(ts->vecs_sensi2p[nadj],-th->time_step,ts->vecs_fpu[nadj]);
400: }
401: if (ts->vecs_fpp) {
402: VecAXPY(ts->vecs_sensi2p[nadj],-th->time_step,ts->vecs_fpp[nadj]);
403: }
404: }
405: }
406: }
408: if (ts->vecs_sensi2) {
409: VecResetArray(ts->vec_sensip_col);
410: MatDenseRestoreColumn(ts->mat_sensip,&xarr);
411: }
412: th->status = TS_STEP_COMPLETE;
413: return(0);
414: }
416: static PetscErrorCode TSAdjointStep_Theta(TS ts)
417: {
418: TS_Theta *th = (TS_Theta*)ts->data;
419: TS quadts = ts->quadraturets;
420: Vec *VecsDeltaLam = th->VecsDeltaLam,*VecsDeltaMu = th->VecsDeltaMu,*VecsSensiTemp = th->VecsSensiTemp;
421: Vec *VecsDeltaLam2 = th->VecsDeltaLam2,*VecsDeltaMu2 = th->VecsDeltaMu2,*VecsSensi2Temp = th->VecsSensi2Temp;
422: PetscInt nadj;
423: Mat J,Jpre,quadJ = NULL,quadJp = NULL;
424: KSP ksp;
425: PetscScalar *xarr;
429: if (th->Theta == 1.) {
430: TSAdjointStepBEuler_Private(ts);
431: return(0);
432: }
433: th->status = TS_STEP_INCOMPLETE;
434: SNESGetKSP(ts->snes,&ksp);
435: TSGetIJacobian(ts,&J,&Jpre,NULL,NULL);
436: if (quadts) {
437: TSGetRHSJacobian(quadts,&quadJ,NULL,NULL,NULL);
438: TSGetRHSJacobianP(quadts,&quadJp,NULL,NULL);
439: }
440: /* If endpoint=1, th->ptime and th->X0 will be used; if endpoint=0, th->stage_time and th->X will be used. */
441: th->stage_time = th->endpoint ? ts->ptime : (ts->ptime+(1.-th->Theta)*ts->time_step); /* time_step is negative*/
442: th->ptime = ts->ptime + ts->time_step;
443: th->time_step = -ts->time_step;
445: /* Build RHS for first-order adjoint */
446: /* Cost function has an integral term */
447: if (quadts) {
448: if (th->endpoint) {
449: TSComputeRHSJacobian(quadts,th->stage_time,ts->vec_sol,quadJ,NULL);
450: } else {
451: TSComputeRHSJacobian(quadts,th->stage_time,th->X,quadJ,NULL);
452: }
453: }
455: for (nadj=0; nadj<ts->numcost; nadj++) {
456: VecCopy(ts->vecs_sensi[nadj],VecsSensiTemp[nadj]);
457: VecScale(VecsSensiTemp[nadj],1./(th->Theta*th->time_step));
458: if (quadJ) {
459: MatDenseGetColumn(quadJ,nadj,&xarr);
460: VecPlaceArray(ts->vec_drdu_col,xarr);
461: VecAXPY(VecsSensiTemp[nadj],1.,ts->vec_drdu_col);
462: VecResetArray(ts->vec_drdu_col);
463: MatDenseRestoreColumn(quadJ,&xarr);
464: }
465: }
467: /* Build LHS for first-order adjoint */
468: th->shift = 1./(th->Theta*th->time_step);
469: if (th->endpoint) {
470: TSComputeSNESJacobian(ts,ts->vec_sol,J,Jpre);
471: } else {
472: TSComputeSNESJacobian(ts,th->X,J,Jpre);
473: }
474: KSPSetOperators(ksp,J,Jpre);
476: /* Solve stage equation LHS*lambda_s = RHS for first-order adjoint */
477: for (nadj=0; nadj<ts->numcost; nadj++) {
478: KSPConvergedReason kspreason;
479: KSPSolveTranspose(ksp,VecsSensiTemp[nadj],VecsDeltaLam[nadj]);
480: KSPGetConvergedReason(ksp,&kspreason);
481: if (kspreason < 0) {
482: ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
483: PetscInfo2(ts,"Step=%D, %Dth cost function, transposed linear solve fails, stopping 1st-order adjoint solve\n",ts->steps,nadj);
484: }
485: }
487: /* Second-order adjoint */
488: if (ts->vecs_sensi2) { /* U_{n+1} */
489: if (!th->endpoint) SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_SUP,"Operation not implemented in TS_Theta");
490: /* Get w1 at t_{n+1} from TLM matrix */
491: MatDenseGetColumn(ts->mat_sensip,0,&xarr);
492: VecPlaceArray(ts->vec_sensip_col,xarr);
493: /* lambda_s^T F_UU w_1 */
494: TSComputeIHessianProductFunctionUU(ts,th->stage_time,ts->vec_sol,VecsDeltaLam,ts->vec_sensip_col,ts->vecs_fuu);
495: VecResetArray(ts->vec_sensip_col);
496: MatDenseRestoreColumn(ts->mat_sensip,&xarr);
497: /* lambda_s^T F_UP w_2 */
498: TSComputeIHessianProductFunctionUP(ts,th->stage_time,ts->vec_sol,VecsDeltaLam,ts->vec_dir,ts->vecs_fup);
499: for (nadj=0; nadj<ts->numcost; nadj++) { /* compute the residual */
500: VecCopy(ts->vecs_sensi2[nadj],VecsSensi2Temp[nadj]);
501: VecScale(VecsSensi2Temp[nadj],th->shift);
502: VecAXPY(VecsSensi2Temp[nadj],-1.,ts->vecs_fuu[nadj]);
503: if (ts->vecs_fup) {
504: VecAXPY(VecsSensi2Temp[nadj],-1.,ts->vecs_fup[nadj]);
505: }
506: }
507: /* Solve stage equation LHS X = RHS for second-order adjoint */
508: for (nadj=0; nadj<ts->numcost; nadj++) {
509: KSPConvergedReason kspreason;
510: KSPSolveTranspose(ksp,VecsSensi2Temp[nadj],VecsDeltaLam2[nadj]);
511: KSPGetConvergedReason(ksp,&kspreason);
512: if (kspreason < 0) {
513: ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
514: PetscInfo2(ts,"Step=%D, %Dth cost function, transposed linear solve fails, stopping 2nd-order adjoint solve\n",ts->steps,nadj);
515: }
516: }
517: }
519: /* Update sensitivities, and evaluate integrals if there is any */
520: if(th->endpoint) { /* two-stage Theta methods with th->Theta!=1, th->Theta==1 leads to BEuler */
521: th->shift = 1./((th->Theta-1.)*th->time_step);
522: th->stage_time = th->ptime;
523: TSComputeSNESJacobian(ts,th->X0,J,Jpre);
524: KSPSetOperators(ksp,J,Jpre);
525: /* R_U at t_n */
526: if (quadts) {
527: TSComputeRHSJacobian(quadts,th->ptime,th->X0,quadJ,NULL);
528: }
529: for (nadj=0; nadj<ts->numcost; nadj++) {
530: MatMultTranspose(J,VecsDeltaLam[nadj],ts->vecs_sensi[nadj]);
531: if (quadJ) {
532: MatDenseGetColumn(quadJ,nadj,&xarr);
533: VecPlaceArray(ts->vec_drdu_col,xarr);
534: VecAXPY(ts->vecs_sensi[nadj],-1.,ts->vec_drdu_col);
535: VecResetArray(ts->vec_drdu_col);
536: MatDenseRestoreColumn(quadJ,&xarr);
537: }
538: VecScale(ts->vecs_sensi[nadj],1./th->shift);
539: }
541: /* Second-order adjoint */
542: if (ts->vecs_sensi2) { /* U_n */
543: /* Get w1 at t_n from TLM matrix */
544: MatDenseGetColumn(th->MatFwdSensip0,0,&xarr);
545: VecPlaceArray(ts->vec_sensip_col,xarr);
546: /* lambda_s^T F_UU w_1 */
547: TSComputeIHessianProductFunctionUU(ts,th->ptime,th->X0,VecsDeltaLam,ts->vec_sensip_col,ts->vecs_fuu);
548: VecResetArray(ts->vec_sensip_col);
549: MatDenseRestoreColumn(th->MatFwdSensip0,&xarr);
550: /* lambda_s^T F_UU w_2 */
551: TSComputeIHessianProductFunctionUP(ts,th->ptime,th->X0,VecsDeltaLam,ts->vec_dir,ts->vecs_fup);
552: for (nadj=0; nadj<ts->numcost; nadj++) {
553: /* 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 */
554: MatMultTranspose(J,VecsDeltaLam2[nadj],ts->vecs_sensi2[nadj]);
555: VecAXPY(ts->vecs_sensi2[nadj],1.,ts->vecs_fuu[nadj]);
556: if (ts->vecs_fup) {
557: VecAXPY(ts->vecs_sensi2[nadj],1.,ts->vecs_fup[nadj]);
558: }
559: VecScale(ts->vecs_sensi2[nadj],1./th->shift);
560: }
561: }
563: th->stage_time = ts->ptime; /* recover the old value */
565: if (ts->vecs_sensip) { /* sensitivities wrt parameters */
566: /* U_{n+1} */
567: th->shift = -1./(th->Theta*th->time_step);
568: TSComputeIJacobianP(ts,th->stage_time,ts->vec_sol,th->Xdot,th->shift,ts->Jacp,PETSC_FALSE);
569: if (quadts) {
570: TSComputeRHSJacobianP(quadts,th->stage_time,ts->vec_sol,quadJp);
571: }
572: for (nadj=0; nadj<ts->numcost; nadj++) {
573: MatMultTranspose(ts->Jacp,VecsDeltaLam[nadj],VecsDeltaMu[nadj]);
574: VecAXPY(ts->vecs_sensip[nadj],-th->time_step*th->Theta,VecsDeltaMu[nadj]);
575: }
576: if (ts->vecs_sensi2p) { /* second-order */
577: /* Get w1 at t_{n+1} from TLM matrix */
578: MatDenseGetColumn(ts->mat_sensip,0,&xarr);
579: VecPlaceArray(ts->vec_sensip_col,xarr);
580: /* lambda_s^T F_PU w_1 */
581: TSComputeIHessianProductFunctionPU(ts,th->stage_time,ts->vec_sol,VecsDeltaLam,ts->vec_sensip_col,ts->vecs_fpu);
582: VecResetArray(ts->vec_sensip_col);
583: MatDenseRestoreColumn(ts->mat_sensip,&xarr);
585: /* lambda_s^T F_PP w_2 */
586: TSComputeIHessianProductFunctionPP(ts,th->stage_time,ts->vec_sol,VecsDeltaLam,ts->vec_dir,ts->vecs_fpp);
587: for (nadj=0; nadj<ts->numcost; nadj++) {
588: /* 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) */
589: MatMultTranspose(ts->Jacp,VecsDeltaLam2[nadj],VecsDeltaMu2[nadj]);
590: VecAXPY(ts->vecs_sensi2p[nadj],-th->time_step*th->Theta,VecsDeltaMu2[nadj]);
591: if (ts->vecs_fpu) {
592: VecAXPY(ts->vecs_sensi2p[nadj],-th->time_step*th->Theta,ts->vecs_fpu[nadj]);
593: }
594: if (ts->vecs_fpp) {
595: VecAXPY(ts->vecs_sensi2p[nadj],-th->time_step*th->Theta,ts->vecs_fpp[nadj]);
596: }
597: }
598: }
600: /* U_s */
601: th->shift = 1./((th->Theta-1.0)*th->time_step);
602: TSComputeIJacobianP(ts,th->ptime,th->X0,th->Xdot,th->shift,ts->Jacp,PETSC_FALSE);
603: if (quadts) {
604: TSComputeRHSJacobianP(quadts,th->ptime,th->X0,quadJp);
605: }
606: for (nadj=0; nadj<ts->numcost; nadj++) {
607: MatMultTranspose(ts->Jacp,VecsDeltaLam[nadj],VecsDeltaMu[nadj]);
608: VecAXPY(ts->vecs_sensip[nadj],-th->time_step*(1.0-th->Theta),VecsDeltaMu[nadj]);
609: if (ts->vecs_sensi2p) { /* second-order */
610: /* Get w1 at t_n from TLM matrix */
611: MatDenseGetColumn(th->MatFwdSensip0,0,&xarr);
612: VecPlaceArray(ts->vec_sensip_col,xarr);
613: /* lambda_s^T F_PU w_1 */
614: TSComputeIHessianProductFunctionPU(ts,th->ptime,th->X0,VecsDeltaLam,ts->vec_sensip_col,ts->vecs_fpu);
615: VecResetArray(ts->vec_sensip_col);
616: MatDenseRestoreColumn(th->MatFwdSensip0,&xarr);
617: /* lambda_s^T F_PP w_2 */
618: TSComputeIHessianProductFunctionPP(ts,th->ptime,th->X0,VecsDeltaLam,ts->vec_dir,ts->vecs_fpp);
619: for (nadj=0; nadj<ts->numcost; nadj++) {
620: /* Mu2 <- Mu2 + h(1-theta) F_P^T Lambda_s + h(1-theta) (lambda_s^T F_UU w_1 + lambda_s^T F_UP w_2) */
621: MatMultTranspose(ts->Jacp,VecsDeltaLam2[nadj],VecsDeltaMu2[nadj]);
622: VecAXPY(ts->vecs_sensi2p[nadj],-th->time_step*(1.0-th->Theta),VecsDeltaMu2[nadj]);
623: if (ts->vecs_fpu) {
624: VecAXPY(ts->vecs_sensi2p[nadj],-th->time_step*(1.0-th->Theta),ts->vecs_fpu[nadj]);
625: }
626: if (ts->vecs_fpp) {
627: VecAXPY(ts->vecs_sensi2p[nadj],-th->time_step*(1.0-th->Theta),ts->vecs_fpp[nadj]);
628: }
629: }
630: }
631: }
632: }
633: } else { /* one-stage case */
634: th->shift = 0.0;
635: TSComputeSNESJacobian(ts,th->X,J,Jpre); /* get -f_y */
636: KSPSetOperators(ksp,J,Jpre);
637: if (quadts) {
638: TSComputeRHSJacobian(quadts,th->stage_time,th->X,quadJ,NULL);
639: }
640: for (nadj=0; nadj<ts->numcost; nadj++) {
641: MatMultTranspose(J,VecsDeltaLam[nadj],VecsSensiTemp[nadj]);
642: VecAXPY(ts->vecs_sensi[nadj],-th->time_step,VecsSensiTemp[nadj]);
643: if (quadJ) {
644: MatDenseGetColumn(quadJ,nadj,&xarr);
645: VecPlaceArray(ts->vec_drdu_col,xarr);
646: VecAXPY(ts->vecs_sensi[nadj],th->time_step,ts->vec_drdu_col);
647: VecResetArray(ts->vec_drdu_col);
648: MatDenseRestoreColumn(quadJ,&xarr);
649: }
650: }
651: if (ts->vecs_sensip) {
652: TSComputeIJacobianP(ts,th->stage_time,th->X,th->Xdot,th->shift,ts->Jacp,PETSC_FALSE);
653: if (quadts) {
654: TSComputeRHSJacobianP(quadts,th->stage_time,th->X,quadJp);
655: }
656: for (nadj=0; nadj<ts->numcost; nadj++) {
657: MatMultTranspose(ts->Jacp,VecsDeltaLam[nadj],VecsDeltaMu[nadj]);
658: VecAXPY(ts->vecs_sensip[nadj],-th->time_step,VecsDeltaMu[nadj]);
659: if (quadJp) {
660: MatDenseGetColumn(quadJp,nadj,&xarr);
661: VecPlaceArray(ts->vec_drdp_col,xarr);
662: VecAXPY(ts->vecs_sensip[nadj],th->time_step,ts->vec_drdp_col);
663: VecResetArray(ts->vec_drdp_col);
664: MatDenseRestoreColumn(quadJp,&xarr);
665: }
666: }
667: }
668: }
670: th->status = TS_STEP_COMPLETE;
671: return(0);
672: }
674: static PetscErrorCode TSInterpolate_Theta(TS ts,PetscReal t,Vec X)
675: {
676: TS_Theta *th = (TS_Theta*)ts->data;
677: PetscReal dt = t - ts->ptime;
681: VecCopy(ts->vec_sol,th->X);
682: if (th->endpoint) dt *= th->Theta;
683: VecWAXPY(X,dt,th->Xdot,th->X);
684: return(0);
685: }
687: static PetscErrorCode TSEvaluateWLTE_Theta(TS ts,NormType wnormtype,PetscInt *order,PetscReal *wlte)
688: {
689: TS_Theta *th = (TS_Theta*)ts->data;
690: Vec X = ts->vec_sol; /* X = solution */
691: Vec Y = th->vec_lte_work; /* Y = X + LTE */
692: PetscReal wltea,wlter;
696: if (!th->vec_sol_prev) {*wlte = -1; return(0);}
697: /* Cannot compute LTE in first step or in restart after event */
698: if (ts->steprestart) {*wlte = -1; return(0);}
699: /* Compute LTE using backward differences with non-constant time step */
700: {
701: PetscReal h = ts->time_step, h_prev = ts->ptime - ts->ptime_prev;
702: PetscReal a = 1 + h_prev/h;
703: PetscScalar scal[3]; Vec vecs[3];
704: scal[0] = +1/a; scal[1] = -1/(a-1); scal[2] = +1/(a*(a-1));
705: vecs[0] = X; vecs[1] = th->X0; vecs[2] = th->vec_sol_prev;
706: VecCopy(X,Y);
707: VecMAXPY(Y,3,scal,vecs);
708: TSErrorWeightedNorm(ts,X,Y,wnormtype,wlte,&wltea,&wlter);
709: }
710: if (order) *order = 2;
711: return(0);
712: }
714: static PetscErrorCode TSRollBack_Theta(TS ts)
715: {
716: TS_Theta *th = (TS_Theta*)ts->data;
717: TS quadts = ts->quadraturets;
721: VecCopy(th->X0,ts->vec_sol);
722: if (quadts && ts->costintegralfwd) {
723: VecCopy(th->VecCostIntegral0,quadts->vec_sol);
724: }
725: th->status = TS_STEP_INCOMPLETE;
726: if (ts->mat_sensip) {
727: MatCopy(th->MatFwdSensip0,ts->mat_sensip,SAME_NONZERO_PATTERN);
728: }
729: if (quadts && quadts->mat_sensip) {
730: MatCopy(th->MatIntegralSensip0,quadts->mat_sensip,SAME_NONZERO_PATTERN);
731: }
732: return(0);
733: }
735: static PetscErrorCode TSForwardStep_Theta(TS ts)
736: {
737: TS_Theta *th = (TS_Theta*)ts->data;
738: TS quadts = ts->quadraturets;
739: Mat MatDeltaFwdSensip = th->MatDeltaFwdSensip;
740: Vec VecDeltaFwdSensipCol = th->VecDeltaFwdSensipCol;
741: PetscInt ntlm;
742: KSP ksp;
743: Mat J,Jpre,quadJ = NULL,quadJp = NULL;
744: PetscScalar *barr,*xarr;
748: MatCopy(ts->mat_sensip,th->MatFwdSensip0,SAME_NONZERO_PATTERN);
750: if (quadts && quadts->mat_sensip) {
751: MatCopy(quadts->mat_sensip,th->MatIntegralSensip0,SAME_NONZERO_PATTERN);
752: }
753: SNESGetKSP(ts->snes,&ksp);
754: TSGetIJacobian(ts,&J,&Jpre,NULL,NULL);
755: if (quadts) {
756: TSGetRHSJacobian(quadts,&quadJ,NULL,NULL,NULL);
757: TSGetRHSJacobianP(quadts,&quadJp,NULL,NULL);
758: }
760: /* Build RHS */
761: if (th->endpoint) { /* 2-stage method*/
762: th->shift = 1./((th->Theta-1.)*th->time_step);
763: TSComputeIJacobian(ts,th->ptime,th->X0,th->Xdot,th->shift,J,Jpre,PETSC_FALSE);
764: MatMatMult(J,ts->mat_sensip,MAT_REUSE_MATRIX,PETSC_DEFAULT,&MatDeltaFwdSensip);
765: MatScale(MatDeltaFwdSensip,(th->Theta-1.)/th->Theta);
767: /* Add the f_p forcing terms */
768: if (ts->Jacp) {
769: TSComputeIJacobianP(ts,th->ptime,th->X0,th->Xdot,th->shift,ts->Jacp,PETSC_FALSE);
770: MatAXPY(MatDeltaFwdSensip,(th->Theta-1.)/th->Theta,ts->Jacp,SUBSET_NONZERO_PATTERN);
771: TSComputeIJacobianP(ts,th->stage_time,ts->vec_sol,th->Xdot,th->shift,ts->Jacp,PETSC_FALSE);
772: MatAXPY(MatDeltaFwdSensip,-1.,ts->Jacp,SUBSET_NONZERO_PATTERN);
773: }
774: } else { /* 1-stage method */
775: th->shift = 0.0;
776: TSComputeIJacobian(ts,th->stage_time,th->X,th->Xdot,th->shift,J,Jpre,PETSC_FALSE);
777: MatMatMult(J,ts->mat_sensip,MAT_REUSE_MATRIX,PETSC_DEFAULT,&MatDeltaFwdSensip);
778: MatScale(MatDeltaFwdSensip,-1.);
780: /* Add the f_p forcing terms */
781: if (ts->Jacp) {
782: TSComputeIJacobianP(ts,th->stage_time,th->X,th->Xdot,th->shift,ts->Jacp,PETSC_FALSE);
783: MatAXPY(MatDeltaFwdSensip,-1.,ts->Jacp,SUBSET_NONZERO_PATTERN);
784: }
785: }
787: /* Build LHS */
788: th->shift = 1/(th->Theta*th->time_step);
789: if (th->endpoint) {
790: TSComputeIJacobian(ts,th->stage_time,ts->vec_sol,th->Xdot,th->shift,J,Jpre,PETSC_FALSE);
791: } else {
792: TSComputeIJacobian(ts,th->stage_time,th->X,th->Xdot,th->shift,J,Jpre,PETSC_FALSE);
793: }
794: KSPSetOperators(ksp,J,Jpre);
796: /*
797: Evaluate the first stage of integral gradients with the 2-stage method:
798: drdu|t_n*S(t_n) + drdp|t_n
799: This is done before the linear solve because the sensitivity variable S(t_n) will be propagated to S(t_{n+1})
800: */
801: if (th->endpoint) { /* 2-stage method only */
802: if (quadts && quadts->mat_sensip) {
803: TSComputeRHSJacobian(quadts,th->ptime,th->X0,quadJ,NULL);
804: TSComputeRHSJacobianP(quadts,th->ptime,th->X0,quadJp);
805: MatTransposeMatMult(ts->mat_sensip,quadJ,MAT_REUSE_MATRIX,PETSC_DEFAULT,&th->MatIntegralSensipTemp);
806: MatAXPY(th->MatIntegralSensipTemp,1,quadJp,SAME_NONZERO_PATTERN);
807: MatAXPY(quadts->mat_sensip,th->time_step*(1.-th->Theta),th->MatIntegralSensipTemp,SAME_NONZERO_PATTERN);
808: }
809: }
811: /* Solve the tangent linear equation for forward sensitivities to parameters */
812: for (ntlm=0; ntlm<th->num_tlm; ntlm++) {
813: KSPConvergedReason kspreason;
814: MatDenseGetColumn(MatDeltaFwdSensip,ntlm,&barr);
815: VecPlaceArray(VecDeltaFwdSensipCol,barr);
816: if (th->endpoint) {
817: MatDenseGetColumn(ts->mat_sensip,ntlm,&xarr);
818: VecPlaceArray(ts->vec_sensip_col,xarr);
819: KSPSolve(ksp,VecDeltaFwdSensipCol,ts->vec_sensip_col);
820: VecResetArray(ts->vec_sensip_col);
821: MatDenseRestoreColumn(ts->mat_sensip,&xarr);
822: } else {
823: KSPSolve(ksp,VecDeltaFwdSensipCol,VecDeltaFwdSensipCol);
824: }
825: KSPGetConvergedReason(ksp,&kspreason);
826: if (kspreason < 0) {
827: ts->reason = TSFORWARD_DIVERGED_LINEAR_SOLVE;
828: PetscInfo2(ts,"Step=%D, %Dth tangent linear solve, linear solve fails, stopping tangent linear solve\n",ts->steps,ntlm);
829: }
830: VecResetArray(VecDeltaFwdSensipCol);
831: MatDenseRestoreColumn(MatDeltaFwdSensip,&barr);
832: }
835: /*
836: Evaluate the second stage of integral gradients with the 2-stage method:
837: drdu|t_{n+1}*S(t_{n+1}) + drdp|t_{n+1}
838: */
839: if (quadts && quadts->mat_sensip) {
840: if (!th->endpoint) {
841: MatAXPY(ts->mat_sensip,1,MatDeltaFwdSensip,SAME_NONZERO_PATTERN); /* stage sensitivity */
842: TSComputeRHSJacobian(quadts,th->stage_time,th->X,quadJ,NULL);
843: TSComputeRHSJacobianP(quadts,th->stage_time,th->X,quadJp);
844: MatTransposeMatMult(ts->mat_sensip,quadJ,MAT_REUSE_MATRIX,PETSC_DEFAULT,&th->MatIntegralSensipTemp);
845: MatAXPY(th->MatIntegralSensipTemp,1,quadJp,SAME_NONZERO_PATTERN);
846: MatAXPY(quadts->mat_sensip,th->time_step,th->MatIntegralSensipTemp,SAME_NONZERO_PATTERN);
847: MatAXPY(ts->mat_sensip,(1.-th->Theta)/th->Theta,MatDeltaFwdSensip,SAME_NONZERO_PATTERN);
848: } else {
849: TSComputeRHSJacobian(quadts,th->stage_time,ts->vec_sol,quadJ,NULL);
850: TSComputeRHSJacobianP(quadts,th->stage_time,ts->vec_sol,quadJp);
851: MatTransposeMatMult(ts->mat_sensip,quadJ,MAT_REUSE_MATRIX,PETSC_DEFAULT,&th->MatIntegralSensipTemp);
852: MatAXPY(th->MatIntegralSensipTemp,1,quadJp,SAME_NONZERO_PATTERN);
853: MatAXPY(quadts->mat_sensip,th->time_step*th->Theta,th->MatIntegralSensipTemp,SAME_NONZERO_PATTERN);
854: }
855: } else {
856: if (!th->endpoint) {
857: MatAXPY(ts->mat_sensip,1./th->Theta,MatDeltaFwdSensip,SAME_NONZERO_PATTERN);
858: }
859: }
860: return(0);
861: }
863: static PetscErrorCode TSForwardGetStages_Theta(TS ts,PetscInt *ns,Mat **stagesensip)
864: {
865: TS_Theta *th = (TS_Theta*)ts->data;
868: if (ns) *ns = 1;
869: if (stagesensip) *stagesensip = th->endpoint ? &(th->MatFwdSensip0) : &(th->MatDeltaFwdSensip);
870: return(0);
871: }
873: /*------------------------------------------------------------*/
874: static PetscErrorCode TSReset_Theta(TS ts)
875: {
876: TS_Theta *th = (TS_Theta*)ts->data;
880: VecDestroy(&th->X);
881: VecDestroy(&th->Xdot);
882: VecDestroy(&th->X0);
883: VecDestroy(&th->affine);
885: VecDestroy(&th->vec_sol_prev);
886: VecDestroy(&th->vec_lte_work);
888: VecDestroy(&th->VecCostIntegral0);
889: return(0);
890: }
892: static PetscErrorCode TSAdjointReset_Theta(TS ts)
893: {
894: TS_Theta *th = (TS_Theta*)ts->data;
898: VecDestroyVecs(ts->numcost,&th->VecsDeltaLam);
899: VecDestroyVecs(ts->numcost,&th->VecsDeltaMu);
900: VecDestroyVecs(ts->numcost,&th->VecsDeltaLam2);
901: VecDestroyVecs(ts->numcost,&th->VecsDeltaMu2);
902: VecDestroyVecs(ts->numcost,&th->VecsSensiTemp);
903: VecDestroyVecs(ts->numcost,&th->VecsSensi2Temp);
904: return(0);
905: }
907: static PetscErrorCode TSDestroy_Theta(TS ts)
908: {
912: TSReset_Theta(ts);
913: if (ts->dm) {
914: DMCoarsenHookRemove(ts->dm,DMCoarsenHook_TSTheta,DMRestrictHook_TSTheta,ts);
915: DMSubDomainHookRemove(ts->dm,DMSubDomainHook_TSTheta,DMSubDomainRestrictHook_TSTheta,ts);
916: }
917: PetscFree(ts->data);
918: PetscObjectComposeFunction((PetscObject)ts,"TSThetaGetTheta_C",NULL);
919: PetscObjectComposeFunction((PetscObject)ts,"TSThetaSetTheta_C",NULL);
920: PetscObjectComposeFunction((PetscObject)ts,"TSThetaGetEndpoint_C",NULL);
921: PetscObjectComposeFunction((PetscObject)ts,"TSThetaSetEndpoint_C",NULL);
922: return(0);
923: }
925: /*
926: This defines the nonlinear equation that is to be solved with SNES
927: G(U) = F[t0+Theta*dt, U, (U-U0)*shift] = 0
928: */
929: static PetscErrorCode SNESTSFormFunction_Theta(SNES snes,Vec x,Vec y,TS ts)
930: {
931: TS_Theta *th = (TS_Theta*)ts->data;
933: Vec X0,Xdot;
934: DM dm,dmsave;
935: PetscReal shift = th->shift;
938: SNESGetDM(snes,&dm);
939: /* When using the endpoint variant, this is actually 1/Theta * Xdot */
940: TSThetaGetX0AndXdot(ts,dm,&X0,&Xdot);
941: if (x != X0) {
942: VecAXPBYPCZ(Xdot,-shift,shift,0,X0,x);
943: } else {
944: VecZeroEntries(Xdot);
945: }
946: /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */
947: dmsave = ts->dm;
948: ts->dm = dm;
949: TSComputeIFunction(ts,th->stage_time,x,Xdot,y,PETSC_FALSE);
950: ts->dm = dmsave;
951: TSThetaRestoreX0AndXdot(ts,dm,&X0,&Xdot);
952: return(0);
953: }
955: static PetscErrorCode SNESTSFormJacobian_Theta(SNES snes,Vec x,Mat A,Mat B,TS ts)
956: {
957: TS_Theta *th = (TS_Theta*)ts->data;
959: Vec Xdot;
960: DM dm,dmsave;
961: PetscReal shift = th->shift;
964: SNESGetDM(snes,&dm);
965: /* Xdot has already been computed in SNESTSFormFunction_Theta (SNES guarantees this) */
966: TSThetaGetX0AndXdot(ts,dm,NULL,&Xdot);
968: dmsave = ts->dm;
969: ts->dm = dm;
970: TSComputeIJacobian(ts,th->stage_time,x,Xdot,shift,A,B,PETSC_FALSE);
971: ts->dm = dmsave;
972: TSThetaRestoreX0AndXdot(ts,dm,NULL,&Xdot);
973: return(0);
974: }
976: static PetscErrorCode TSForwardSetUp_Theta(TS ts)
977: {
978: TS_Theta *th = (TS_Theta*)ts->data;
979: TS quadts = ts->quadraturets;
983: /* combine sensitivities to parameters and sensitivities to initial values into one array */
984: th->num_tlm = ts->num_parameters;
985: MatDuplicate(ts->mat_sensip,MAT_DO_NOT_COPY_VALUES,&th->MatDeltaFwdSensip);
986: if (quadts && quadts->mat_sensip) {
987: MatDuplicate(quadts->mat_sensip,MAT_DO_NOT_COPY_VALUES,&th->MatIntegralSensipTemp);
988: MatDuplicate(quadts->mat_sensip,MAT_DO_NOT_COPY_VALUES,&th->MatIntegralSensip0);
989: }
990: /* backup sensitivity results for roll-backs */
991: MatDuplicate(ts->mat_sensip,MAT_DO_NOT_COPY_VALUES,&th->MatFwdSensip0);
993: VecDuplicate(ts->vec_sol,&th->VecDeltaFwdSensipCol);
994: return(0);
995: }
997: static PetscErrorCode TSForwardReset_Theta(TS ts)
998: {
999: TS_Theta *th = (TS_Theta*)ts->data;
1000: TS quadts = ts->quadraturets;
1004: if (quadts && quadts->mat_sensip) {
1005: MatDestroy(&th->MatIntegralSensipTemp);
1006: MatDestroy(&th->MatIntegralSensip0);
1007: }
1008: VecDestroy(&th->VecDeltaFwdSensipCol);
1009: MatDestroy(&th->MatDeltaFwdSensip);
1010: MatDestroy(&th->MatFwdSensip0);
1011: return(0);
1012: }
1014: static PetscErrorCode TSSetUp_Theta(TS ts)
1015: {
1016: TS_Theta *th = (TS_Theta*)ts->data;
1017: TS quadts = ts->quadraturets;
1018: PetscBool match;
1022: if (!th->VecCostIntegral0 && quadts && ts->costintegralfwd) { /* back up cost integral */
1023: VecDuplicate(quadts->vec_sol,&th->VecCostIntegral0);
1024: }
1025: if (!th->X) {
1026: VecDuplicate(ts->vec_sol,&th->X);
1027: }
1028: if (!th->Xdot) {
1029: VecDuplicate(ts->vec_sol,&th->Xdot);
1030: }
1031: if (!th->X0) {
1032: VecDuplicate(ts->vec_sol,&th->X0);
1033: }
1034: if (th->endpoint) {
1035: VecDuplicate(ts->vec_sol,&th->affine);
1036: }
1038: th->order = (th->Theta == 0.5) ? 2 : 1;
1040: TSGetDM(ts,&ts->dm);
1041: DMCoarsenHookAdd(ts->dm,DMCoarsenHook_TSTheta,DMRestrictHook_TSTheta,ts);
1042: DMSubDomainHookAdd(ts->dm,DMSubDomainHook_TSTheta,DMSubDomainRestrictHook_TSTheta,ts);
1044: TSGetAdapt(ts,&ts->adapt);
1045: TSAdaptCandidatesClear(ts->adapt);
1046: PetscObjectTypeCompare((PetscObject)ts->adapt,TSADAPTNONE,&match);
1047: if (!match) {
1048: VecDuplicate(ts->vec_sol,&th->vec_sol_prev);
1049: VecDuplicate(ts->vec_sol,&th->vec_lte_work);
1050: }
1051: TSGetSNES(ts,&ts->snes);
1052: return(0);
1053: }
1055: /*------------------------------------------------------------*/
1057: static PetscErrorCode TSAdjointSetUp_Theta(TS ts)
1058: {
1059: TS_Theta *th = (TS_Theta*)ts->data;
1063: VecDuplicateVecs(ts->vecs_sensi[0],ts->numcost,&th->VecsDeltaLam);
1064: VecDuplicateVecs(ts->vecs_sensi[0],ts->numcost,&th->VecsSensiTemp);
1065: if (ts->vecs_sensip) {
1066: VecDuplicateVecs(ts->vecs_sensip[0],ts->numcost,&th->VecsDeltaMu);
1067: }
1068: if (ts->vecs_sensi2) {
1069: VecDuplicateVecs(ts->vecs_sensi[0],ts->numcost,&th->VecsDeltaLam2);
1070: VecDuplicateVecs(ts->vecs_sensi2[0],ts->numcost,&th->VecsSensi2Temp);
1071: /* hack ts to make implicit TS solver work when users provide only explicit versions of callbacks (RHSFunction,RHSJacobian,RHSHessian etc.) */
1072: if (!ts->ihessianproduct_fuu) ts->vecs_fuu = ts->vecs_guu;
1073: if (!ts->ihessianproduct_fup) ts->vecs_fup = ts->vecs_gup;
1074: }
1075: if (ts->vecs_sensi2p) {
1076: VecDuplicateVecs(ts->vecs_sensi2p[0],ts->numcost,&th->VecsDeltaMu2);
1077: /* hack ts to make implicit TS solver work when users provide only explicit versions of callbacks (RHSFunction,RHSJacobian,RHSHessian etc.) */
1078: if (!ts->ihessianproduct_fpu) ts->vecs_fpu = ts->vecs_gpu;
1079: if (!ts->ihessianproduct_fpp) ts->vecs_fpp = ts->vecs_gpp;
1080: }
1081: return(0);
1082: }
1084: static PetscErrorCode TSSetFromOptions_Theta(PetscOptionItems *PetscOptionsObject,TS ts)
1085: {
1086: TS_Theta *th = (TS_Theta*)ts->data;
1090: PetscOptionsHead(PetscOptionsObject,"Theta ODE solver options");
1091: {
1092: PetscOptionsReal("-ts_theta_theta","Location of stage (0<Theta<=1)","TSThetaSetTheta",th->Theta,&th->Theta,NULL);
1093: PetscOptionsBool("-ts_theta_endpoint","Use the endpoint instead of midpoint form of the Theta method","TSThetaSetEndpoint",th->endpoint,&th->endpoint,NULL);
1094: PetscOptionsBool("-ts_theta_initial_guess_extrapolate","Extrapolate stage initial guess from previous solution (sometimes unstable)","TSThetaSetExtrapolate",th->extrapolate,&th->extrapolate,NULL);
1095: }
1096: PetscOptionsTail();
1097: return(0);
1098: }
1100: static PetscErrorCode TSView_Theta(TS ts,PetscViewer viewer)
1101: {
1102: TS_Theta *th = (TS_Theta*)ts->data;
1103: PetscBool iascii;
1107: PetscObjectTypeCompare((PetscObject)viewer,PETSCVIEWERASCII,&iascii);
1108: if (iascii) {
1109: PetscViewerASCIIPrintf(viewer," Theta=%g\n",(double)th->Theta);
1110: PetscViewerASCIIPrintf(viewer," Extrapolation=%s\n",th->extrapolate ? "yes" : "no");
1111: }
1112: return(0);
1113: }
1115: static PetscErrorCode TSThetaGetTheta_Theta(TS ts,PetscReal *theta)
1116: {
1117: TS_Theta *th = (TS_Theta*)ts->data;
1120: *theta = th->Theta;
1121: return(0);
1122: }
1124: static PetscErrorCode TSThetaSetTheta_Theta(TS ts,PetscReal theta)
1125: {
1126: TS_Theta *th = (TS_Theta*)ts->data;
1129: if (theta <= 0 || 1 < theta) SETERRQ1(PetscObjectComm((PetscObject)ts),PETSC_ERR_ARG_OUTOFRANGE,"Theta %g not in range (0,1]",(double)theta);
1130: th->Theta = theta;
1131: th->order = (th->Theta == 0.5) ? 2 : 1;
1132: return(0);
1133: }
1135: static PetscErrorCode TSThetaGetEndpoint_Theta(TS ts,PetscBool *endpoint)
1136: {
1137: TS_Theta *th = (TS_Theta*)ts->data;
1140: *endpoint = th->endpoint;
1141: return(0);
1142: }
1144: static PetscErrorCode TSThetaSetEndpoint_Theta(TS ts,PetscBool flg)
1145: {
1146: TS_Theta *th = (TS_Theta*)ts->data;
1149: th->endpoint = flg;
1150: return(0);
1151: }
1153: #if defined(PETSC_HAVE_COMPLEX)
1154: static PetscErrorCode TSComputeLinearStability_Theta(TS ts,PetscReal xr,PetscReal xi,PetscReal *yr,PetscReal *yi)
1155: {
1156: PetscComplex z = xr + xi*PETSC_i,f;
1157: TS_Theta *th = (TS_Theta*)ts->data;
1158: const PetscReal one = 1.0;
1161: f = (one + (one - th->Theta)*z)/(one - th->Theta*z);
1162: *yr = PetscRealPartComplex(f);
1163: *yi = PetscImaginaryPartComplex(f);
1164: return(0);
1165: }
1166: #endif
1168: static PetscErrorCode TSGetStages_Theta(TS ts,PetscInt *ns,Vec **Y)
1169: {
1170: TS_Theta *th = (TS_Theta*)ts->data;
1173: if (ns) *ns = 1;
1174: if (Y) *Y = th->endpoint ? &(th->X0) : &(th->X);
1175: return(0);
1176: }
1178: /* ------------------------------------------------------------ */
1179: /*MC
1180: TSTHETA - DAE solver using the implicit Theta method
1182: Level: beginner
1184: Options Database:
1185: + -ts_theta_theta <Theta> - Location of stage (0<Theta<=1)
1186: . -ts_theta_endpoint <flag> - Use the endpoint (like Crank-Nicholson) instead of midpoint form of the Theta method
1187: - -ts_theta_initial_guess_extrapolate <flg> - Extrapolate stage initial guess from previous solution (sometimes unstable)
1189: Notes:
1190: $ -ts_type theta -ts_theta_theta 1.0 corresponds to backward Euler (TSBEULER)
1191: $ -ts_type theta -ts_theta_theta 0.5 corresponds to the implicit midpoint rule
1192: $ -ts_type theta -ts_theta_theta 0.5 -ts_theta_endpoint corresponds to Crank-Nicholson (TSCN)
1194: This method can be applied to DAE.
1196: This method is cast as a 1-stage implicit Runge-Kutta method.
1198: .vb
1199: Theta | Theta
1200: -------------
1201: | 1
1202: .ve
1204: For the default Theta=0.5, this is also known as the implicit midpoint rule.
1206: When the endpoint variant is chosen, the method becomes a 2-stage method with first stage explicit:
1208: .vb
1209: 0 | 0 0
1210: 1 | 1-Theta Theta
1211: -------------------
1212: | 1-Theta Theta
1213: .ve
1215: For the default Theta=0.5, this is the trapezoid rule (also known as Crank-Nicolson, see TSCN).
1217: To apply a diagonally implicit RK method to DAE, the stage formula
1219: $ Y_i = X + h sum_j a_ij Y'_j
1221: is interpreted as a formula for Y'_i in terms of Y_i and known values (Y'_j, j<i)
1223: .seealso: TSCreate(), TS, TSSetType(), TSCN, TSBEULER, TSThetaSetTheta(), TSThetaSetEndpoint()
1225: M*/
1226: PETSC_EXTERN PetscErrorCode TSCreate_Theta(TS ts)
1227: {
1228: TS_Theta *th;
1232: ts->ops->reset = TSReset_Theta;
1233: ts->ops->adjointreset = TSAdjointReset_Theta;
1234: ts->ops->destroy = TSDestroy_Theta;
1235: ts->ops->view = TSView_Theta;
1236: ts->ops->setup = TSSetUp_Theta;
1237: ts->ops->adjointsetup = TSAdjointSetUp_Theta;
1238: ts->ops->adjointreset = TSAdjointReset_Theta;
1239: ts->ops->step = TSStep_Theta;
1240: ts->ops->interpolate = TSInterpolate_Theta;
1241: ts->ops->evaluatewlte = TSEvaluateWLTE_Theta;
1242: ts->ops->rollback = TSRollBack_Theta;
1243: ts->ops->setfromoptions = TSSetFromOptions_Theta;
1244: ts->ops->snesfunction = SNESTSFormFunction_Theta;
1245: ts->ops->snesjacobian = SNESTSFormJacobian_Theta;
1246: #if defined(PETSC_HAVE_COMPLEX)
1247: ts->ops->linearstability = TSComputeLinearStability_Theta;
1248: #endif
1249: ts->ops->getstages = TSGetStages_Theta;
1250: ts->ops->adjointstep = TSAdjointStep_Theta;
1251: ts->ops->adjointintegral = TSAdjointCostIntegral_Theta;
1252: ts->ops->forwardintegral = TSForwardCostIntegral_Theta;
1253: ts->default_adapt_type = TSADAPTNONE;
1255: ts->ops->forwardsetup = TSForwardSetUp_Theta;
1256: ts->ops->forwardreset = TSForwardReset_Theta;
1257: ts->ops->forwardstep = TSForwardStep_Theta;
1258: ts->ops->forwardgetstages = TSForwardGetStages_Theta;
1260: ts->usessnes = PETSC_TRUE;
1262: PetscNewLog(ts,&th);
1263: ts->data = (void*)th;
1265: th->VecsDeltaLam = NULL;
1266: th->VecsDeltaMu = NULL;
1267: th->VecsSensiTemp = NULL;
1268: th->VecsSensi2Temp = NULL;
1270: th->extrapolate = PETSC_FALSE;
1271: th->Theta = 0.5;
1272: th->order = 2;
1273: PetscObjectComposeFunction((PetscObject)ts,"TSThetaGetTheta_C",TSThetaGetTheta_Theta);
1274: PetscObjectComposeFunction((PetscObject)ts,"TSThetaSetTheta_C",TSThetaSetTheta_Theta);
1275: PetscObjectComposeFunction((PetscObject)ts,"TSThetaGetEndpoint_C",TSThetaGetEndpoint_Theta);
1276: PetscObjectComposeFunction((PetscObject)ts,"TSThetaSetEndpoint_C",TSThetaSetEndpoint_Theta);
1277: return(0);
1278: }
1280: /*@
1281: TSThetaGetTheta - Get the abscissa of the stage in (0,1].
1283: Not Collective
1285: Input Parameter:
1286: . ts - timestepping context
1288: Output Parameter:
1289: . theta - stage abscissa
1291: Note:
1292: Use of this function is normally only required to hack TSTHETA to use a modified integration scheme.
1294: Level: Advanced
1296: .seealso: TSThetaSetTheta()
1297: @*/
1298: PetscErrorCode TSThetaGetTheta(TS ts,PetscReal *theta)
1299: {
1305: PetscUseMethod(ts,"TSThetaGetTheta_C",(TS,PetscReal*),(ts,theta));
1306: return(0);
1307: }
1309: /*@
1310: TSThetaSetTheta - Set the abscissa of the stage in (0,1].
1312: Not Collective
1314: Input Parameter:
1315: + ts - timestepping context
1316: - theta - stage abscissa
1318: Options Database:
1319: . -ts_theta_theta <theta>
1321: Level: Intermediate
1323: .seealso: TSThetaGetTheta()
1324: @*/
1325: PetscErrorCode TSThetaSetTheta(TS ts,PetscReal theta)
1326: {
1331: PetscTryMethod(ts,"TSThetaSetTheta_C",(TS,PetscReal),(ts,theta));
1332: return(0);
1333: }
1335: /*@
1336: TSThetaGetEndpoint - Gets whether to use the endpoint variant of the method (e.g. trapezoid/Crank-Nicolson instead of midpoint rule).
1338: Not Collective
1340: Input Parameter:
1341: . ts - timestepping context
1343: Output Parameter:
1344: . endpoint - PETSC_TRUE when using the endpoint variant
1346: Level: Advanced
1348: .seealso: TSThetaSetEndpoint(), TSTHETA, TSCN
1349: @*/
1350: PetscErrorCode TSThetaGetEndpoint(TS ts,PetscBool *endpoint)
1351: {
1357: PetscUseMethod(ts,"TSThetaGetEndpoint_C",(TS,PetscBool*),(ts,endpoint));
1358: return(0);
1359: }
1361: /*@
1362: TSThetaSetEndpoint - Sets whether to use the endpoint variant of the method (e.g. trapezoid/Crank-Nicolson instead of midpoint rule).
1364: Not Collective
1366: Input Parameter:
1367: + ts - timestepping context
1368: - flg - PETSC_TRUE to use the endpoint variant
1370: Options Database:
1371: . -ts_theta_endpoint <flg>
1373: Level: Intermediate
1375: .seealso: TSTHETA, TSCN
1376: @*/
1377: PetscErrorCode TSThetaSetEndpoint(TS ts,PetscBool flg)
1378: {
1383: PetscTryMethod(ts,"TSThetaSetEndpoint_C",(TS,PetscBool),(ts,flg));
1384: return(0);
1385: }
1387: /*
1388: * TSBEULER and TSCN are straightforward specializations of TSTHETA.
1389: * The creation functions for these specializations are below.
1390: */
1392: static PetscErrorCode TSSetUp_BEuler(TS ts)
1393: {
1394: TS_Theta *th = (TS_Theta*)ts->data;
1398: if (th->Theta != 1.0) SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_OPT_OVERWRITE,"Can not change the default value (1) of theta when using backward Euler\n");
1399: if (th->endpoint) SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_OPT_OVERWRITE,"Can not change to the endpoint form of the Theta methods when using backward Euler\n");
1400: TSSetUp_Theta(ts);
1401: return(0);
1402: }
1404: static PetscErrorCode TSView_BEuler(TS ts,PetscViewer viewer)
1405: {
1407: return(0);
1408: }
1410: /*MC
1411: TSBEULER - ODE solver using the implicit backward Euler method
1413: Level: beginner
1415: Notes:
1416: TSBEULER is equivalent to TSTHETA with Theta=1.0
1418: $ -ts_type theta -ts_theta_theta 1.0
1420: .seealso: TSCreate(), TS, TSSetType(), TSEULER, TSCN, TSTHETA
1422: M*/
1423: PETSC_EXTERN PetscErrorCode TSCreate_BEuler(TS ts)
1424: {
1428: TSCreate_Theta(ts);
1429: TSThetaSetTheta(ts,1.0);
1430: TSThetaSetEndpoint(ts,PETSC_FALSE);
1431: ts->ops->setup = TSSetUp_BEuler;
1432: ts->ops->view = TSView_BEuler;
1433: return(0);
1434: }
1436: static PetscErrorCode TSSetUp_CN(TS ts)
1437: {
1438: TS_Theta *th = (TS_Theta*)ts->data;
1442: if (th->Theta != 0.5) SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_OPT_OVERWRITE,"Can not change the default value (0.5) of theta when using Crank-Nicolson\n");
1443: if (!th->endpoint) SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_OPT_OVERWRITE,"Can not change to the midpoint form of the Theta methods when using Crank-Nicolson\n");
1444: TSSetUp_Theta(ts);
1445: return(0);
1446: }
1448: static PetscErrorCode TSView_CN(TS ts,PetscViewer viewer)
1449: {
1451: return(0);
1452: }
1454: /*MC
1455: TSCN - ODE solver using the implicit Crank-Nicolson method.
1457: Level: beginner
1459: Notes:
1460: TSCN is equivalent to TSTHETA with Theta=0.5 and the "endpoint" option set. I.e.
1462: $ -ts_type theta -ts_theta_theta 0.5 -ts_theta_endpoint
1464: .seealso: TSCreate(), TS, TSSetType(), TSBEULER, TSTHETA
1466: M*/
1467: PETSC_EXTERN PetscErrorCode TSCreate_CN(TS ts)
1468: {
1472: TSCreate_Theta(ts);
1473: TSThetaSetTheta(ts,0.5);
1474: TSThetaSetEndpoint(ts,PETSC_TRUE);
1475: ts->ops->setup = TSSetUp_CN;
1476: ts->ops->view = TSView_CN;
1477: return(0);
1478: }