Actual source code: theta.c
petsc-3.13.1 2020-05-02
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 shift; /* Shift parameter for SNES Jacobian, used by forward, TLM and adjoint */
16: PetscInt order;
17: PetscBool endpoint;
18: PetscBool extrapolate;
19: TSStepStatus status;
20: Vec VecCostIntegral0; /* Backup for roll-backs due to events, used by cost integral */
21: PetscReal ptime0; /* Backup for ts->ptime, the start time of current time step, used by TLM and cost integral */
22: PetscReal time_step0; /* Backup for ts->timestep, the step size of current time step, used by TLM and cost integral*/
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->ptime0,th->X0,ts->vec_costintegrand);
142: VecAXPY(quadts->vec_sol,th->time_step0*(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_step0*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_step0,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: {
168: TS_Theta *th = (TS_Theta*)ts->data;
172: /* Like TSForwardCostIntegral(), the adjoint cost integral evaluation relies on ptime0 and time_step0. */
173: th->ptime0 = ts->ptime + ts->time_step;
174: th->time_step0 = -ts->time_step;
175: TSThetaEvaluateCostIntegral(ts);
176: return(0);
177: }
179: static PetscErrorCode TSTheta_SNESSolve(TS ts,Vec b,Vec x)
180: {
181: PetscInt nits,lits;
185: SNESSolve(ts->snes,b,x);
186: SNESGetIterationNumber(ts->snes,&nits);
187: SNESGetLinearSolveIterations(ts->snes,&lits);
188: ts->snes_its += nits; ts->ksp_its += lits;
189: return(0);
190: }
192: static PetscErrorCode TSStep_Theta(TS ts)
193: {
194: TS_Theta *th = (TS_Theta*)ts->data;
195: PetscInt rejections = 0;
196: PetscBool stageok,accept = PETSC_TRUE;
197: PetscReal next_time_step = ts->time_step;
201: if (!ts->steprollback) {
202: if (th->vec_sol_prev) { VecCopy(th->X0,th->vec_sol_prev); }
203: VecCopy(ts->vec_sol,th->X0);
204: }
206: th->status = TS_STEP_INCOMPLETE;
207: while (!ts->reason && th->status != TS_STEP_COMPLETE) {
208: th->shift = 1/(th->Theta*ts->time_step);
209: th->stage_time = ts->ptime + (th->endpoint ? (PetscReal)1 : th->Theta)*ts->time_step;
210: VecCopy(th->X0,th->X);
211: if (th->extrapolate && !ts->steprestart) {
212: VecAXPY(th->X,1/th->shift,th->Xdot);
213: }
214: if (th->endpoint) { /* This formulation assumes linear time-independent mass matrix */
215: if (!th->affine) {VecDuplicate(ts->vec_sol,&th->affine);}
216: VecZeroEntries(th->Xdot);
217: TSComputeIFunction(ts,ts->ptime,th->X0,th->Xdot,th->affine,PETSC_FALSE);
218: VecScale(th->affine,(th->Theta-1)/th->Theta);
219: } else if (th->affine) { /* Just in case th->endpoint is changed between calls to TSStep_Theta() */
220: VecZeroEntries(th->affine);
221: }
222: TSPreStage(ts,th->stage_time);
223: TSTheta_SNESSolve(ts,th->affine,th->X);
224: TSPostStage(ts,th->stage_time,0,&th->X);
225: TSAdaptCheckStage(ts->adapt,ts,th->stage_time,th->X,&stageok);
226: if (!stageok) goto reject_step;
228: th->status = TS_STEP_PENDING;
229: if (th->endpoint) {
230: VecCopy(th->X,ts->vec_sol);
231: } else {
232: VecAXPBYPCZ(th->Xdot,-th->shift,th->shift,0,th->X0,th->X);
233: VecAXPY(ts->vec_sol,ts->time_step,th->Xdot);
234: }
235: TSAdaptChoose(ts->adapt,ts,ts->time_step,NULL,&next_time_step,&accept);
236: th->status = accept ? TS_STEP_COMPLETE : TS_STEP_INCOMPLETE;
237: if (!accept) {
238: VecCopy(th->X0,ts->vec_sol);
239: ts->time_step = next_time_step;
240: goto reject_step;
241: }
243: if (ts->forward_solve || ts->costintegralfwd) { /* Save the info for the later use in cost integral evaluation */
244: th->ptime0 = ts->ptime;
245: th->time_step0 = ts->time_step;
246: }
247: ts->ptime += ts->time_step;
248: ts->time_step = next_time_step;
249: break;
251: reject_step:
252: ts->reject++; accept = PETSC_FALSE;
253: if (!ts->reason && ++rejections > ts->max_reject && ts->max_reject >= 0) {
254: ts->reason = TS_DIVERGED_STEP_REJECTED;
255: PetscInfo2(ts,"Step=%D, step rejections %D greater than current TS allowed, stopping solve\n",ts->steps,rejections);
256: }
257: }
258: return(0);
259: }
261: static PetscErrorCode TSAdjointStepBEuler_Private(TS ts)
262: {
263: TS_Theta *th = (TS_Theta*)ts->data;
264: TS quadts = ts->quadraturets;
265: Vec *VecsDeltaLam = th->VecsDeltaLam,*VecsDeltaMu = th->VecsDeltaMu,*VecsSensiTemp = th->VecsSensiTemp;
266: Vec *VecsDeltaLam2 = th->VecsDeltaLam2,*VecsDeltaMu2 = th->VecsDeltaMu2,*VecsSensi2Temp = th->VecsSensi2Temp;
267: PetscInt nadj;
268: Mat J,Jpre,quadJ = NULL,quadJp = NULL;
269: KSP ksp;
270: PetscScalar *xarr;
271: TSEquationType eqtype;
272: PetscBool isexplicitode = PETSC_FALSE;
273: PetscReal adjoint_time_step;
277: TSGetEquationType(ts,&eqtype);
278: if (eqtype == TS_EQ_ODE_EXPLICIT) {
279: isexplicitode = PETSC_TRUE;
280: VecsDeltaLam = ts->vecs_sensi;
281: VecsDeltaLam2 = ts->vecs_sensi2;
282: }
283: th->status = TS_STEP_INCOMPLETE;
284: SNESGetKSP(ts->snes,&ksp);
285: TSGetIJacobian(ts,&J,&Jpre,NULL,NULL);
286: if (quadts) {
287: TSGetRHSJacobian(quadts,&quadJ,NULL,NULL,NULL);
288: TSGetRHSJacobianP(quadts,&quadJp,NULL,NULL);
289: }
291: th->stage_time = ts->ptime;
292: adjoint_time_step = -ts->time_step; /* always positive since time_step is negative */
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./adjoint_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./adjoint_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],adjoint_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],adjoint_time_step);
368: VecCopy(VecsSensi2Temp[nadj],ts->vecs_sensi2[nadj]);
369: }
370: }
371: }
372: if (ts->vecs_sensip) {
373: TSComputeIJacobianP(ts,th->stage_time,th->X,th->Xdot,1./adjoint_time_step,ts->Jacp,PETSC_FALSE); /* get -f_p */
374: if (quadts) {
375: TSComputeRHSJacobianP(quadts,th->stage_time,th->X,quadJp);
376: }
377: if (ts->vecs_sensi2p) {
378: /* lambda_s^T F_PU w_1 */
379: TSComputeIHessianProductFunctionPU(ts,th->stage_time,th->X,VecsDeltaLam,ts->vec_sensip_col,ts->vecs_fpu);
380: /* lambda_s^T F_PP w_2 */
381: TSComputeIHessianProductFunctionPP(ts,th->stage_time,th->X,VecsDeltaLam,ts->vec_dir,ts->vecs_fpp);
382: }
384: for (nadj=0; nadj<ts->numcost; nadj++) {
385: MatMultTranspose(ts->Jacp,VecsDeltaLam[nadj],VecsDeltaMu[nadj]);
386: VecAXPY(ts->vecs_sensip[nadj],-adjoint_time_step,VecsDeltaMu[nadj]);
387: if (quadJp) {
388: MatDenseGetColumn(quadJp,nadj,&xarr);
389: VecPlaceArray(ts->vec_drdp_col,xarr);
390: VecAXPY(ts->vecs_sensip[nadj],adjoint_time_step,ts->vec_drdp_col);
391: VecResetArray(ts->vec_drdp_col);
392: MatDenseRestoreColumn(quadJp,&xarr);
393: }
394: if (ts->vecs_sensi2p) {
395: MatMultTranspose(ts->Jacp,VecsDeltaLam2[nadj],VecsDeltaMu2[nadj]);
396: VecAXPY(ts->vecs_sensi2p[nadj],-adjoint_time_step,VecsDeltaMu2[nadj]);
397: if (ts->vecs_fpu) {
398: VecAXPY(ts->vecs_sensi2p[nadj],-adjoint_time_step,ts->vecs_fpu[nadj]);
399: }
400: if (ts->vecs_fpp) {
401: VecAXPY(ts->vecs_sensi2p[nadj],-adjoint_time_step,ts->vecs_fpp[nadj]);
402: }
403: }
404: }
405: }
407: if (ts->vecs_sensi2) {
408: VecResetArray(ts->vec_sensip_col);
409: MatDenseRestoreColumn(ts->mat_sensip,&xarr);
410: }
411: th->status = TS_STEP_COMPLETE;
412: return(0);
413: }
415: static PetscErrorCode TSAdjointStep_Theta(TS ts)
416: {
417: TS_Theta *th = (TS_Theta*)ts->data;
418: TS quadts = ts->quadraturets;
419: Vec *VecsDeltaLam = th->VecsDeltaLam,*VecsDeltaMu = th->VecsDeltaMu,*VecsSensiTemp = th->VecsSensiTemp;
420: Vec *VecsDeltaLam2 = th->VecsDeltaLam2,*VecsDeltaMu2 = th->VecsDeltaMu2,*VecsSensi2Temp = th->VecsSensi2Temp;
421: PetscInt nadj;
422: Mat J,Jpre,quadJ = NULL,quadJp = NULL;
423: KSP ksp;
424: PetscScalar *xarr;
425: PetscReal adjoint_time_step;
426: PetscReal adjoint_ptime; /* end time of the adjoint time step (ts->ptime is the start time, ususally ts->ptime is larger than adjoint_ptime) */
430: if (th->Theta == 1.) {
431: TSAdjointStepBEuler_Private(ts);
432: return(0);
433: }
434: th->status = TS_STEP_INCOMPLETE;
435: SNESGetKSP(ts->snes,&ksp);
436: TSGetIJacobian(ts,&J,&Jpre,NULL,NULL);
437: if (quadts) {
438: TSGetRHSJacobian(quadts,&quadJ,NULL,NULL,NULL);
439: TSGetRHSJacobianP(quadts,&quadJp,NULL,NULL);
440: }
441: /* If endpoint=1, th->ptime and th->X0 will be used; if endpoint=0, th->stage_time and th->X will be used. */
442: th->stage_time = th->endpoint ? ts->ptime : (ts->ptime+(1.-th->Theta)*ts->time_step);
443: adjoint_ptime = ts->ptime + ts->time_step;
444: adjoint_time_step = -ts->time_step; /* always positive since time_step is negative */
446: /* Build RHS for first-order adjoint */
447: /* Cost function has an integral term */
448: if (quadts) {
449: if (th->endpoint) {
450: TSComputeRHSJacobian(quadts,th->stage_time,ts->vec_sol,quadJ,NULL);
451: } else {
452: TSComputeRHSJacobian(quadts,th->stage_time,th->X,quadJ,NULL);
453: }
454: }
456: for (nadj=0; nadj<ts->numcost; nadj++) {
457: VecCopy(ts->vecs_sensi[nadj],VecsSensiTemp[nadj]);
458: VecScale(VecsSensiTemp[nadj],1./(th->Theta*adjoint_time_step));
459: if (quadJ) {
460: MatDenseGetColumn(quadJ,nadj,&xarr);
461: VecPlaceArray(ts->vec_drdu_col,xarr);
462: VecAXPY(VecsSensiTemp[nadj],1.,ts->vec_drdu_col);
463: VecResetArray(ts->vec_drdu_col);
464: MatDenseRestoreColumn(quadJ,&xarr);
465: }
466: }
468: /* Build LHS for first-order adjoint */
469: th->shift = 1./(th->Theta*adjoint_time_step);
470: if (th->endpoint) {
471: TSComputeSNESJacobian(ts,ts->vec_sol,J,Jpre);
472: } else {
473: TSComputeSNESJacobian(ts,th->X,J,Jpre);
474: }
475: KSPSetOperators(ksp,J,Jpre);
477: /* Solve stage equation LHS*lambda_s = RHS for first-order adjoint */
478: for (nadj=0; nadj<ts->numcost; nadj++) {
479: KSPConvergedReason kspreason;
480: KSPSolveTranspose(ksp,VecsSensiTemp[nadj],VecsDeltaLam[nadj]);
481: KSPGetConvergedReason(ksp,&kspreason);
482: if (kspreason < 0) {
483: ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
484: PetscInfo2(ts,"Step=%D, %Dth cost function, transposed linear solve fails, stopping 1st-order adjoint solve\n",ts->steps,nadj);
485: }
486: }
488: /* Second-order adjoint */
489: if (ts->vecs_sensi2) { /* U_{n+1} */
490: if (!th->endpoint) SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_SUP,"Operation not implemented in TS_Theta");
491: /* Get w1 at t_{n+1} from TLM matrix */
492: MatDenseGetColumn(ts->mat_sensip,0,&xarr);
493: VecPlaceArray(ts->vec_sensip_col,xarr);
494: /* lambda_s^T F_UU w_1 */
495: TSComputeIHessianProductFunctionUU(ts,th->stage_time,ts->vec_sol,VecsDeltaLam,ts->vec_sensip_col,ts->vecs_fuu);
496: VecResetArray(ts->vec_sensip_col);
497: MatDenseRestoreColumn(ts->mat_sensip,&xarr);
498: /* lambda_s^T F_UP w_2 */
499: TSComputeIHessianProductFunctionUP(ts,th->stage_time,ts->vec_sol,VecsDeltaLam,ts->vec_dir,ts->vecs_fup);
500: for (nadj=0; nadj<ts->numcost; nadj++) { /* compute the residual */
501: VecCopy(ts->vecs_sensi2[nadj],VecsSensi2Temp[nadj]);
502: VecScale(VecsSensi2Temp[nadj],th->shift);
503: VecAXPY(VecsSensi2Temp[nadj],-1.,ts->vecs_fuu[nadj]);
504: if (ts->vecs_fup) {
505: VecAXPY(VecsSensi2Temp[nadj],-1.,ts->vecs_fup[nadj]);
506: }
507: }
508: /* Solve stage equation LHS X = RHS for second-order adjoint */
509: for (nadj=0; nadj<ts->numcost; nadj++) {
510: KSPConvergedReason kspreason;
511: KSPSolveTranspose(ksp,VecsSensi2Temp[nadj],VecsDeltaLam2[nadj]);
512: KSPGetConvergedReason(ksp,&kspreason);
513: if (kspreason < 0) {
514: ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
515: PetscInfo2(ts,"Step=%D, %Dth cost function, transposed linear solve fails, stopping 2nd-order adjoint solve\n",ts->steps,nadj);
516: }
517: }
518: }
520: /* Update sensitivities, and evaluate integrals if there is any */
521: if(th->endpoint) { /* two-stage Theta methods with th->Theta!=1, th->Theta==1 leads to BEuler */
522: th->shift = 1./((th->Theta-1.)*adjoint_time_step);
523: th->stage_time = adjoint_ptime;
524: TSComputeSNESJacobian(ts,th->X0,J,Jpre);
525: KSPSetOperators(ksp,J,Jpre);
526: /* R_U at t_n */
527: if (quadts) {
528: TSComputeRHSJacobian(quadts,adjoint_ptime,th->X0,quadJ,NULL);
529: }
530: for (nadj=0; nadj<ts->numcost; nadj++) {
531: MatMultTranspose(J,VecsDeltaLam[nadj],ts->vecs_sensi[nadj]);
532: if (quadJ) {
533: MatDenseGetColumn(quadJ,nadj,&xarr);
534: VecPlaceArray(ts->vec_drdu_col,xarr);
535: VecAXPY(ts->vecs_sensi[nadj],-1.,ts->vec_drdu_col);
536: VecResetArray(ts->vec_drdu_col);
537: MatDenseRestoreColumn(quadJ,&xarr);
538: }
539: VecScale(ts->vecs_sensi[nadj],1./th->shift);
540: }
542: /* Second-order adjoint */
543: if (ts->vecs_sensi2) { /* U_n */
544: /* Get w1 at t_n from TLM matrix */
545: MatDenseGetColumn(th->MatFwdSensip0,0,&xarr);
546: VecPlaceArray(ts->vec_sensip_col,xarr);
547: /* lambda_s^T F_UU w_1 */
548: TSComputeIHessianProductFunctionUU(ts,adjoint_ptime,th->X0,VecsDeltaLam,ts->vec_sensip_col,ts->vecs_fuu);
549: VecResetArray(ts->vec_sensip_col);
550: MatDenseRestoreColumn(th->MatFwdSensip0,&xarr);
551: /* lambda_s^T F_UU w_2 */
552: TSComputeIHessianProductFunctionUP(ts,adjoint_ptime,th->X0,VecsDeltaLam,ts->vec_dir,ts->vecs_fup);
553: for (nadj=0; nadj<ts->numcost; nadj++) {
554: /* 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 */
555: MatMultTranspose(J,VecsDeltaLam2[nadj],ts->vecs_sensi2[nadj]);
556: VecAXPY(ts->vecs_sensi2[nadj],1.,ts->vecs_fuu[nadj]);
557: if (ts->vecs_fup) {
558: VecAXPY(ts->vecs_sensi2[nadj],1.,ts->vecs_fup[nadj]);
559: }
560: VecScale(ts->vecs_sensi2[nadj],1./th->shift);
561: }
562: }
564: th->stage_time = ts->ptime; /* recover the old value */
566: if (ts->vecs_sensip) { /* sensitivities wrt parameters */
567: /* U_{n+1} */
568: TSComputeIJacobianP(ts,th->stage_time,ts->vec_sol,th->Xdot,-1./(th->Theta*adjoint_time_step),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],-adjoint_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],-adjoint_time_step*th->Theta,VecsDeltaMu2[nadj]);
591: if (ts->vecs_fpu) {
592: VecAXPY(ts->vecs_sensi2p[nadj],-adjoint_time_step*th->Theta,ts->vecs_fpu[nadj]);
593: }
594: if (ts->vecs_fpp) {
595: VecAXPY(ts->vecs_sensi2p[nadj],-adjoint_time_step*th->Theta,ts->vecs_fpp[nadj]);
596: }
597: }
598: }
600: /* U_s */
601: TSComputeIJacobianP(ts,adjoint_ptime,th->X0,th->Xdot,1./((th->Theta-1.0)*adjoint_time_step),ts->Jacp,PETSC_FALSE);
602: if (quadts) {
603: TSComputeRHSJacobianP(quadts,adjoint_ptime,th->X0,quadJp);
604: }
605: for (nadj=0; nadj<ts->numcost; nadj++) {
606: MatMultTranspose(ts->Jacp,VecsDeltaLam[nadj],VecsDeltaMu[nadj]);
607: VecAXPY(ts->vecs_sensip[nadj],-adjoint_time_step*(1.0-th->Theta),VecsDeltaMu[nadj]);
608: if (ts->vecs_sensi2p) { /* second-order */
609: /* Get w1 at t_n from TLM matrix */
610: MatDenseGetColumn(th->MatFwdSensip0,0,&xarr);
611: VecPlaceArray(ts->vec_sensip_col,xarr);
612: /* lambda_s^T F_PU w_1 */
613: TSComputeIHessianProductFunctionPU(ts,adjoint_ptime,th->X0,VecsDeltaLam,ts->vec_sensip_col,ts->vecs_fpu);
614: VecResetArray(ts->vec_sensip_col);
615: MatDenseRestoreColumn(th->MatFwdSensip0,&xarr);
616: /* lambda_s^T F_PP w_2 */
617: TSComputeIHessianProductFunctionPP(ts,adjoint_ptime,th->X0,VecsDeltaLam,ts->vec_dir,ts->vecs_fpp);
618: for (nadj=0; nadj<ts->numcost; nadj++) {
619: /* Mu2 <- Mu2 + h(1-theta) F_P^T Lambda_s + h(1-theta) (lambda_s^T F_UU w_1 + lambda_s^T F_UP w_2) */
620: MatMultTranspose(ts->Jacp,VecsDeltaLam2[nadj],VecsDeltaMu2[nadj]);
621: VecAXPY(ts->vecs_sensi2p[nadj],-adjoint_time_step*(1.0-th->Theta),VecsDeltaMu2[nadj]);
622: if (ts->vecs_fpu) {
623: VecAXPY(ts->vecs_sensi2p[nadj],-adjoint_time_step*(1.0-th->Theta),ts->vecs_fpu[nadj]);
624: }
625: if (ts->vecs_fpp) {
626: VecAXPY(ts->vecs_sensi2p[nadj],-adjoint_time_step*(1.0-th->Theta),ts->vecs_fpp[nadj]);
627: }
628: }
629: }
630: }
631: }
632: } else { /* one-stage case */
633: th->shift = 0.0;
634: TSComputeSNESJacobian(ts,th->X,J,Jpre); /* get -f_y */
635: KSPSetOperators(ksp,J,Jpre);
636: if (quadts) {
637: TSComputeRHSJacobian(quadts,th->stage_time,th->X,quadJ,NULL);
638: }
639: for (nadj=0; nadj<ts->numcost; nadj++) {
640: MatMultTranspose(J,VecsDeltaLam[nadj],VecsSensiTemp[nadj]);
641: VecAXPY(ts->vecs_sensi[nadj],-adjoint_time_step,VecsSensiTemp[nadj]);
642: if (quadJ) {
643: MatDenseGetColumn(quadJ,nadj,&xarr);
644: VecPlaceArray(ts->vec_drdu_col,xarr);
645: VecAXPY(ts->vecs_sensi[nadj],adjoint_time_step,ts->vec_drdu_col);
646: VecResetArray(ts->vec_drdu_col);
647: MatDenseRestoreColumn(quadJ,&xarr);
648: }
649: }
650: if (ts->vecs_sensip) {
651: TSComputeIJacobianP(ts,th->stage_time,th->X,th->Xdot,th->shift,ts->Jacp,PETSC_FALSE);
652: if (quadts) {
653: TSComputeRHSJacobianP(quadts,th->stage_time,th->X,quadJp);
654: }
655: for (nadj=0; nadj<ts->numcost; nadj++) {
656: MatMultTranspose(ts->Jacp,VecsDeltaLam[nadj],VecsDeltaMu[nadj]);
657: VecAXPY(ts->vecs_sensip[nadj],-adjoint_time_step,VecsDeltaMu[nadj]);
658: if (quadJp) {
659: MatDenseGetColumn(quadJp,nadj,&xarr);
660: VecPlaceArray(ts->vec_drdp_col,xarr);
661: VecAXPY(ts->vecs_sensip[nadj],adjoint_time_step,ts->vec_drdp_col);
662: VecResetArray(ts->vec_drdp_col);
663: MatDenseRestoreColumn(quadJp,&xarr);
664: }
665: }
666: }
667: }
669: th->status = TS_STEP_COMPLETE;
670: return(0);
671: }
673: static PetscErrorCode TSInterpolate_Theta(TS ts,PetscReal t,Vec X)
674: {
675: TS_Theta *th = (TS_Theta*)ts->data;
676: PetscReal dt = t - ts->ptime;
680: VecCopy(ts->vec_sol,th->X);
681: if (th->endpoint) dt *= th->Theta;
682: VecWAXPY(X,dt,th->Xdot,th->X);
683: return(0);
684: }
686: static PetscErrorCode TSEvaluateWLTE_Theta(TS ts,NormType wnormtype,PetscInt *order,PetscReal *wlte)
687: {
688: TS_Theta *th = (TS_Theta*)ts->data;
689: Vec X = ts->vec_sol; /* X = solution */
690: Vec Y = th->vec_lte_work; /* Y = X + LTE */
691: PetscReal wltea,wlter;
695: if (!th->vec_sol_prev) {*wlte = -1; return(0);}
696: /* Cannot compute LTE in first step or in restart after event */
697: if (ts->steprestart) {*wlte = -1; return(0);}
698: /* Compute LTE using backward differences with non-constant time step */
699: {
700: PetscReal h = ts->time_step, h_prev = ts->ptime - ts->ptime_prev;
701: PetscReal a = 1 + h_prev/h;
702: PetscScalar scal[3]; Vec vecs[3];
703: scal[0] = +1/a; scal[1] = -1/(a-1); scal[2] = +1/(a*(a-1));
704: vecs[0] = X; vecs[1] = th->X0; vecs[2] = th->vec_sol_prev;
705: VecCopy(X,Y);
706: VecMAXPY(Y,3,scal,vecs);
707: TSErrorWeightedNorm(ts,X,Y,wnormtype,wlte,&wltea,&wlter);
708: }
709: if (order) *order = 2;
710: return(0);
711: }
713: static PetscErrorCode TSRollBack_Theta(TS ts)
714: {
715: TS_Theta *th = (TS_Theta*)ts->data;
716: TS quadts = ts->quadraturets;
720: VecCopy(th->X0,ts->vec_sol);
721: if (quadts && ts->costintegralfwd) {
722: VecCopy(th->VecCostIntegral0,quadts->vec_sol);
723: }
724: th->status = TS_STEP_INCOMPLETE;
725: if (ts->mat_sensip) {
726: MatCopy(th->MatFwdSensip0,ts->mat_sensip,SAME_NONZERO_PATTERN);
727: }
728: if (quadts && quadts->mat_sensip) {
729: MatCopy(th->MatIntegralSensip0,quadts->mat_sensip,SAME_NONZERO_PATTERN);
730: }
731: return(0);
732: }
734: static PetscErrorCode TSForwardStep_Theta(TS ts)
735: {
736: TS_Theta *th = (TS_Theta*)ts->data;
737: TS quadts = ts->quadraturets;
738: Mat MatDeltaFwdSensip = th->MatDeltaFwdSensip;
739: Vec VecDeltaFwdSensipCol = th->VecDeltaFwdSensipCol;
740: PetscInt ntlm;
741: KSP ksp;
742: Mat J,Jpre,quadJ = NULL,quadJp = NULL;
743: PetscScalar *barr,*xarr;
744: PetscReal previous_shift;
748: previous_shift = th->shift;
749: MatCopy(ts->mat_sensip,th->MatFwdSensip0,SAME_NONZERO_PATTERN);
751: if (quadts && quadts->mat_sensip) {
752: MatCopy(quadts->mat_sensip,th->MatIntegralSensip0,SAME_NONZERO_PATTERN);
753: }
754: SNESGetKSP(ts->snes,&ksp);
755: TSGetIJacobian(ts,&J,&Jpre,NULL,NULL);
756: if (quadts) {
757: TSGetRHSJacobian(quadts,&quadJ,NULL,NULL,NULL);
758: TSGetRHSJacobianP(quadts,&quadJp,NULL,NULL);
759: }
761: /* Build RHS */
762: if (th->endpoint) { /* 2-stage method*/
763: th->shift = 1./((th->Theta-1.)*th->time_step0);
764: TSComputeIJacobian(ts,th->ptime0,th->X0,th->Xdot,th->shift,J,Jpre,PETSC_FALSE);
765: MatMatMult(J,ts->mat_sensip,MAT_REUSE_MATRIX,PETSC_DEFAULT,&MatDeltaFwdSensip);
766: MatScale(MatDeltaFwdSensip,(th->Theta-1.)/th->Theta);
768: /* Add the f_p forcing terms */
769: if (ts->Jacp) {
770: TSComputeIJacobianP(ts,th->ptime0,th->X0,th->Xdot,th->shift,ts->Jacp,PETSC_FALSE);
771: MatAXPY(MatDeltaFwdSensip,(th->Theta-1.)/th->Theta,ts->Jacp,SUBSET_NONZERO_PATTERN);
772: TSComputeIJacobianP(ts,th->stage_time,ts->vec_sol,th->Xdot,th->shift,ts->Jacp,PETSC_FALSE);
773: MatAXPY(MatDeltaFwdSensip,-1.,ts->Jacp,SUBSET_NONZERO_PATTERN);
774: }
775: } else { /* 1-stage method */
776: th->shift = 0.0;
777: TSComputeIJacobian(ts,th->stage_time,th->X,th->Xdot,th->shift,J,Jpre,PETSC_FALSE);
778: MatMatMult(J,ts->mat_sensip,MAT_REUSE_MATRIX,PETSC_DEFAULT,&MatDeltaFwdSensip);
779: MatScale(MatDeltaFwdSensip,-1.);
781: /* Add the f_p forcing terms */
782: if (ts->Jacp) {
783: TSComputeIJacobianP(ts,th->stage_time,th->X,th->Xdot,th->shift,ts->Jacp,PETSC_FALSE);
784: MatAXPY(MatDeltaFwdSensip,-1.,ts->Jacp,SUBSET_NONZERO_PATTERN);
785: }
786: }
788: /* Build LHS */
789: th->shift = previous_shift; /* recover the previous shift used in TSStep_Theta() */
790: if (th->endpoint) {
791: TSComputeIJacobian(ts,th->stage_time,ts->vec_sol,th->Xdot,th->shift,J,Jpre,PETSC_FALSE);
792: } else {
793: TSComputeIJacobian(ts,th->stage_time,th->X,th->Xdot,th->shift,J,Jpre,PETSC_FALSE);
794: }
795: KSPSetOperators(ksp,J,Jpre);
797: /*
798: Evaluate the first stage of integral gradients with the 2-stage method:
799: drdu|t_n*S(t_n) + drdp|t_n
800: This is done before the linear solve because the sensitivity variable S(t_n) will be propagated to S(t_{n+1})
801: */
802: if (th->endpoint) { /* 2-stage method only */
803: if (quadts && quadts->mat_sensip) {
804: TSComputeRHSJacobian(quadts,th->ptime0,th->X0,quadJ,NULL);
805: TSComputeRHSJacobianP(quadts,th->ptime0,th->X0,quadJp);
806: MatTransposeMatMult(ts->mat_sensip,quadJ,MAT_REUSE_MATRIX,PETSC_DEFAULT,&th->MatIntegralSensipTemp);
807: MatAXPY(th->MatIntegralSensipTemp,1,quadJp,SAME_NONZERO_PATTERN);
808: MatAXPY(quadts->mat_sensip,th->time_step0*(1.-th->Theta),th->MatIntegralSensipTemp,SAME_NONZERO_PATTERN);
809: }
810: }
812: /* Solve the tangent linear equation for forward sensitivities to parameters */
813: for (ntlm=0; ntlm<th->num_tlm; ntlm++) {
814: KSPConvergedReason kspreason;
815: MatDenseGetColumn(MatDeltaFwdSensip,ntlm,&barr);
816: VecPlaceArray(VecDeltaFwdSensipCol,barr);
817: if (th->endpoint) {
818: MatDenseGetColumn(ts->mat_sensip,ntlm,&xarr);
819: VecPlaceArray(ts->vec_sensip_col,xarr);
820: KSPSolve(ksp,VecDeltaFwdSensipCol,ts->vec_sensip_col);
821: VecResetArray(ts->vec_sensip_col);
822: MatDenseRestoreColumn(ts->mat_sensip,&xarr);
823: } else {
824: KSPSolve(ksp,VecDeltaFwdSensipCol,VecDeltaFwdSensipCol);
825: }
826: KSPGetConvergedReason(ksp,&kspreason);
827: if (kspreason < 0) {
828: ts->reason = TSFORWARD_DIVERGED_LINEAR_SOLVE;
829: PetscInfo2(ts,"Step=%D, %Dth tangent linear solve, linear solve fails, stopping tangent linear solve\n",ts->steps,ntlm);
830: }
831: VecResetArray(VecDeltaFwdSensipCol);
832: MatDenseRestoreColumn(MatDeltaFwdSensip,&barr);
833: }
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_step0,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_step0*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;
1039: th->shift = 1/(th->Theta*ts->time_step);
1041: TSGetDM(ts,&ts->dm);
1042: DMCoarsenHookAdd(ts->dm,DMCoarsenHook_TSTheta,DMRestrictHook_TSTheta,ts);
1043: DMSubDomainHookAdd(ts->dm,DMSubDomainHook_TSTheta,DMSubDomainRestrictHook_TSTheta,ts);
1045: TSGetAdapt(ts,&ts->adapt);
1046: TSAdaptCandidatesClear(ts->adapt);
1047: PetscObjectTypeCompare((PetscObject)ts->adapt,TSADAPTNONE,&match);
1048: if (!match) {
1049: VecDuplicate(ts->vec_sol,&th->vec_sol_prev);
1050: VecDuplicate(ts->vec_sol,&th->vec_lte_work);
1051: }
1052: TSGetSNES(ts,&ts->snes);
1053: return(0);
1054: }
1056: /*------------------------------------------------------------*/
1058: static PetscErrorCode TSAdjointSetUp_Theta(TS ts)
1059: {
1060: TS_Theta *th = (TS_Theta*)ts->data;
1064: VecDuplicateVecs(ts->vecs_sensi[0],ts->numcost,&th->VecsDeltaLam);
1065: VecDuplicateVecs(ts->vecs_sensi[0],ts->numcost,&th->VecsSensiTemp);
1066: if (ts->vecs_sensip) {
1067: VecDuplicateVecs(ts->vecs_sensip[0],ts->numcost,&th->VecsDeltaMu);
1068: }
1069: if (ts->vecs_sensi2) {
1070: VecDuplicateVecs(ts->vecs_sensi[0],ts->numcost,&th->VecsDeltaLam2);
1071: VecDuplicateVecs(ts->vecs_sensi2[0],ts->numcost,&th->VecsSensi2Temp);
1072: /* hack ts to make implicit TS solver work when users provide only explicit versions of callbacks (RHSFunction,RHSJacobian,RHSHessian etc.) */
1073: if (!ts->ihessianproduct_fuu) ts->vecs_fuu = ts->vecs_guu;
1074: if (!ts->ihessianproduct_fup) ts->vecs_fup = ts->vecs_gup;
1075: }
1076: if (ts->vecs_sensi2p) {
1077: VecDuplicateVecs(ts->vecs_sensi2p[0],ts->numcost,&th->VecsDeltaMu2);
1078: /* hack ts to make implicit TS solver work when users provide only explicit versions of callbacks (RHSFunction,RHSJacobian,RHSHessian etc.) */
1079: if (!ts->ihessianproduct_fpu) ts->vecs_fpu = ts->vecs_gpu;
1080: if (!ts->ihessianproduct_fpp) ts->vecs_fpp = ts->vecs_gpp;
1081: }
1082: return(0);
1083: }
1085: static PetscErrorCode TSSetFromOptions_Theta(PetscOptionItems *PetscOptionsObject,TS ts)
1086: {
1087: TS_Theta *th = (TS_Theta*)ts->data;
1091: PetscOptionsHead(PetscOptionsObject,"Theta ODE solver options");
1092: {
1093: PetscOptionsReal("-ts_theta_theta","Location of stage (0<Theta<=1)","TSThetaSetTheta",th->Theta,&th->Theta,NULL);
1094: PetscOptionsBool("-ts_theta_endpoint","Use the endpoint instead of midpoint form of the Theta method","TSThetaSetEndpoint",th->endpoint,&th->endpoint,NULL);
1095: PetscOptionsBool("-ts_theta_initial_guess_extrapolate","Extrapolate stage initial guess from previous solution (sometimes unstable)","TSThetaSetExtrapolate",th->extrapolate,&th->extrapolate,NULL);
1096: }
1097: PetscOptionsTail();
1098: return(0);
1099: }
1101: static PetscErrorCode TSView_Theta(TS ts,PetscViewer viewer)
1102: {
1103: TS_Theta *th = (TS_Theta*)ts->data;
1104: PetscBool iascii;
1108: PetscObjectTypeCompare((PetscObject)viewer,PETSCVIEWERASCII,&iascii);
1109: if (iascii) {
1110: PetscViewerASCIIPrintf(viewer," Theta=%g\n",(double)th->Theta);
1111: PetscViewerASCIIPrintf(viewer," Extrapolation=%s\n",th->extrapolate ? "yes" : "no");
1112: }
1113: return(0);
1114: }
1116: static PetscErrorCode TSThetaGetTheta_Theta(TS ts,PetscReal *theta)
1117: {
1118: TS_Theta *th = (TS_Theta*)ts->data;
1121: *theta = th->Theta;
1122: return(0);
1123: }
1125: static PetscErrorCode TSThetaSetTheta_Theta(TS ts,PetscReal theta)
1126: {
1127: TS_Theta *th = (TS_Theta*)ts->data;
1130: if (theta <= 0 || 1 < theta) SETERRQ1(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: return(0);
1134: }
1136: static PetscErrorCode TSThetaGetEndpoint_Theta(TS ts,PetscBool *endpoint)
1137: {
1138: TS_Theta *th = (TS_Theta*)ts->data;
1141: *endpoint = th->endpoint;
1142: return(0);
1143: }
1145: static PetscErrorCode TSThetaSetEndpoint_Theta(TS ts,PetscBool flg)
1146: {
1147: TS_Theta *th = (TS_Theta*)ts->data;
1150: th->endpoint = flg;
1151: return(0);
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;
1159: const PetscReal one = 1.0;
1162: f = (one + (one - th->Theta)*z)/(one - th->Theta*z);
1163: *yr = PetscRealPartComplex(f);
1164: *yi = PetscImaginaryPartComplex(f);
1165: return(0);
1166: }
1167: #endif
1169: static PetscErrorCode TSGetStages_Theta(TS ts,PetscInt *ns,Vec **Y)
1170: {
1171: TS_Theta *th = (TS_Theta*)ts->data;
1174: if (ns) *ns = 1;
1175: if (Y) *Y = th->endpoint ? &(th->X0) : &(th->X);
1176: return(0);
1177: }
1179: /* ------------------------------------------------------------ */
1180: /*MC
1181: TSTHETA - DAE solver using the implicit Theta method
1183: Level: beginner
1185: Options Database:
1186: + -ts_theta_theta <Theta> - Location of stage (0<Theta<=1)
1187: . -ts_theta_endpoint <flag> - Use the endpoint (like Crank-Nicholson) instead of midpoint form of the Theta method
1188: - -ts_theta_initial_guess_extrapolate <flg> - Extrapolate stage initial guess from previous solution (sometimes unstable)
1190: Notes:
1191: $ -ts_type theta -ts_theta_theta 1.0 corresponds to backward Euler (TSBEULER)
1192: $ -ts_type theta -ts_theta_theta 0.5 corresponds to the implicit midpoint rule
1193: $ -ts_type theta -ts_theta_theta 0.5 -ts_theta_endpoint corresponds to Crank-Nicholson (TSCN)
1195: This method can be applied to DAE.
1197: This method is cast as a 1-stage implicit Runge-Kutta method.
1199: .vb
1200: Theta | Theta
1201: -------------
1202: | 1
1203: .ve
1205: For the default Theta=0.5, this is also known as the implicit midpoint rule.
1207: When the endpoint variant is chosen, the method becomes a 2-stage method with first stage explicit:
1209: .vb
1210: 0 | 0 0
1211: 1 | 1-Theta Theta
1212: -------------------
1213: | 1-Theta Theta
1214: .ve
1216: For the default Theta=0.5, this is the trapezoid rule (also known as Crank-Nicolson, see TSCN).
1218: To apply a diagonally implicit RK method to DAE, the stage formula
1220: $ Y_i = X + h sum_j a_ij Y'_j
1222: is interpreted as a formula for Y'_i in terms of Y_i and known values (Y'_j, j<i)
1224: .seealso: TSCreate(), TS, TSSetType(), TSCN, TSBEULER, TSThetaSetTheta(), TSThetaSetEndpoint()
1226: M*/
1227: PETSC_EXTERN PetscErrorCode TSCreate_Theta(TS ts)
1228: {
1229: TS_Theta *th;
1233: ts->ops->reset = TSReset_Theta;
1234: ts->ops->adjointreset = TSAdjointReset_Theta;
1235: ts->ops->destroy = TSDestroy_Theta;
1236: ts->ops->view = TSView_Theta;
1237: ts->ops->setup = TSSetUp_Theta;
1238: ts->ops->adjointsetup = TSAdjointSetUp_Theta;
1239: ts->ops->adjointreset = TSAdjointReset_Theta;
1240: ts->ops->step = TSStep_Theta;
1241: ts->ops->interpolate = TSInterpolate_Theta;
1242: ts->ops->evaluatewlte = TSEvaluateWLTE_Theta;
1243: ts->ops->rollback = TSRollBack_Theta;
1244: ts->ops->setfromoptions = TSSetFromOptions_Theta;
1245: ts->ops->snesfunction = SNESTSFormFunction_Theta;
1246: ts->ops->snesjacobian = SNESTSFormJacobian_Theta;
1247: #if defined(PETSC_HAVE_COMPLEX)
1248: ts->ops->linearstability = TSComputeLinearStability_Theta;
1249: #endif
1250: ts->ops->getstages = TSGetStages_Theta;
1251: ts->ops->adjointstep = TSAdjointStep_Theta;
1252: ts->ops->adjointintegral = TSAdjointCostIntegral_Theta;
1253: ts->ops->forwardintegral = TSForwardCostIntegral_Theta;
1254: ts->default_adapt_type = TSADAPTNONE;
1256: ts->ops->forwardsetup = TSForwardSetUp_Theta;
1257: ts->ops->forwardreset = TSForwardReset_Theta;
1258: ts->ops->forwardstep = TSForwardStep_Theta;
1259: ts->ops->forwardgetstages = TSForwardGetStages_Theta;
1261: ts->usessnes = PETSC_TRUE;
1263: PetscNewLog(ts,&th);
1264: ts->data = (void*)th;
1266: th->VecsDeltaLam = NULL;
1267: th->VecsDeltaMu = NULL;
1268: th->VecsSensiTemp = NULL;
1269: th->VecsSensi2Temp = NULL;
1271: th->extrapolate = PETSC_FALSE;
1272: th->Theta = 0.5;
1273: th->order = 2;
1274: PetscObjectComposeFunction((PetscObject)ts,"TSThetaGetTheta_C",TSThetaGetTheta_Theta);
1275: PetscObjectComposeFunction((PetscObject)ts,"TSThetaSetTheta_C",TSThetaSetTheta_Theta);
1276: PetscObjectComposeFunction((PetscObject)ts,"TSThetaGetEndpoint_C",TSThetaGetEndpoint_Theta);
1277: PetscObjectComposeFunction((PetscObject)ts,"TSThetaSetEndpoint_C",TSThetaSetEndpoint_Theta);
1278: return(0);
1279: }
1281: /*@
1282: TSThetaGetTheta - Get the abscissa of the stage in (0,1].
1284: Not Collective
1286: Input Parameter:
1287: . ts - timestepping context
1289: Output Parameter:
1290: . theta - stage abscissa
1292: Note:
1293: Use of this function is normally only required to hack TSTHETA to use a modified integration scheme.
1295: Level: Advanced
1297: .seealso: TSThetaSetTheta()
1298: @*/
1299: PetscErrorCode TSThetaGetTheta(TS ts,PetscReal *theta)
1300: {
1306: PetscUseMethod(ts,"TSThetaGetTheta_C",(TS,PetscReal*),(ts,theta));
1307: return(0);
1308: }
1310: /*@
1311: TSThetaSetTheta - Set the abscissa of the stage in (0,1].
1313: Not Collective
1315: Input Parameter:
1316: + ts - timestepping context
1317: - theta - stage abscissa
1319: Options Database:
1320: . -ts_theta_theta <theta>
1322: Level: Intermediate
1324: .seealso: TSThetaGetTheta()
1325: @*/
1326: PetscErrorCode TSThetaSetTheta(TS ts,PetscReal theta)
1327: {
1332: PetscTryMethod(ts,"TSThetaSetTheta_C",(TS,PetscReal),(ts,theta));
1333: return(0);
1334: }
1336: /*@
1337: TSThetaGetEndpoint - Gets whether to use the endpoint variant of the method (e.g. trapezoid/Crank-Nicolson instead of midpoint rule).
1339: Not Collective
1341: Input Parameter:
1342: . ts - timestepping context
1344: Output Parameter:
1345: . endpoint - PETSC_TRUE when using the endpoint variant
1347: Level: Advanced
1349: .seealso: TSThetaSetEndpoint(), TSTHETA, TSCN
1350: @*/
1351: PetscErrorCode TSThetaGetEndpoint(TS ts,PetscBool *endpoint)
1352: {
1358: PetscUseMethod(ts,"TSThetaGetEndpoint_C",(TS,PetscBool*),(ts,endpoint));
1359: return(0);
1360: }
1362: /*@
1363: TSThetaSetEndpoint - Sets whether to use the endpoint variant of the method (e.g. trapezoid/Crank-Nicolson instead of midpoint rule).
1365: Not Collective
1367: Input Parameter:
1368: + ts - timestepping context
1369: - flg - PETSC_TRUE to use the endpoint variant
1371: Options Database:
1372: . -ts_theta_endpoint <flg>
1374: Level: Intermediate
1376: .seealso: TSTHETA, TSCN
1377: @*/
1378: PetscErrorCode TSThetaSetEndpoint(TS ts,PetscBool flg)
1379: {
1384: PetscTryMethod(ts,"TSThetaSetEndpoint_C",(TS,PetscBool),(ts,flg));
1385: return(0);
1386: }
1388: /*
1389: * TSBEULER and TSCN are straightforward specializations of TSTHETA.
1390: * The creation functions for these specializations are below.
1391: */
1393: static PetscErrorCode TSSetUp_BEuler(TS ts)
1394: {
1395: TS_Theta *th = (TS_Theta*)ts->data;
1399: 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");
1400: 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");
1401: TSSetUp_Theta(ts);
1402: return(0);
1403: }
1405: static PetscErrorCode TSView_BEuler(TS ts,PetscViewer viewer)
1406: {
1408: return(0);
1409: }
1411: /*MC
1412: TSBEULER - ODE solver using the implicit backward Euler method
1414: Level: beginner
1416: Notes:
1417: TSBEULER is equivalent to TSTHETA with Theta=1.0
1419: $ -ts_type theta -ts_theta_theta 1.0
1421: .seealso: TSCreate(), TS, TSSetType(), TSEULER, TSCN, TSTHETA
1423: M*/
1424: PETSC_EXTERN PetscErrorCode TSCreate_BEuler(TS ts)
1425: {
1429: TSCreate_Theta(ts);
1430: TSThetaSetTheta(ts,1.0);
1431: TSThetaSetEndpoint(ts,PETSC_FALSE);
1432: ts->ops->setup = TSSetUp_BEuler;
1433: ts->ops->view = TSView_BEuler;
1434: return(0);
1435: }
1437: static PetscErrorCode TSSetUp_CN(TS ts)
1438: {
1439: TS_Theta *th = (TS_Theta*)ts->data;
1443: 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");
1444: 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");
1445: TSSetUp_Theta(ts);
1446: return(0);
1447: }
1449: static PetscErrorCode TSView_CN(TS ts,PetscViewer viewer)
1450: {
1452: return(0);
1453: }
1455: /*MC
1456: TSCN - ODE solver using the implicit Crank-Nicolson method.
1458: Level: beginner
1460: Notes:
1461: TSCN is equivalent to TSTHETA with Theta=0.5 and the "endpoint" option set. I.e.
1463: $ -ts_type theta -ts_theta_theta 0.5 -ts_theta_endpoint
1465: .seealso: TSCreate(), TS, TSSetType(), TSBEULER, TSTHETA
1467: M*/
1468: PETSC_EXTERN PetscErrorCode TSCreate_CN(TS ts)
1469: {
1473: TSCreate_Theta(ts);
1474: TSThetaSetTheta(ts,0.5);
1475: TSThetaSetEndpoint(ts,PETSC_TRUE);
1476: ts->ops->setup = TSSetUp_CN;
1477: ts->ops->view = TSView_CN;
1478: return(0);
1479: }