Actual source code: theta.c
petsc-3.13.2 2020-06-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: th->shift = 1./adjoint_time_step;
313: TSComputeSNESJacobian(ts,th->X,J,Jpre);
314: KSPSetOperators(ksp,J,Jpre);
316: /* Solve stage equation LHS*lambda_s = RHS for first-order adjoint */
317: for (nadj=0; nadj<ts->numcost; nadj++) {
318: KSPConvergedReason kspreason;
319: KSPSolveTranspose(ksp,VecsSensiTemp[nadj],VecsDeltaLam[nadj]);
320: KSPGetConvergedReason(ksp,&kspreason);
321: if (kspreason < 0) {
322: ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
323: PetscInfo2(ts,"Step=%D, %Dth cost function, transposed linear solve fails, stopping 1st-order adjoint solve\n",ts->steps,nadj);
324: }
325: }
327: if (ts->vecs_sensi2) { /* U_{n+1} */
328: /* Get w1 at t_{n+1} from TLM matrix */
329: MatDenseGetColumn(ts->mat_sensip,0,&xarr);
330: VecPlaceArray(ts->vec_sensip_col,xarr);
331: /* lambda_s^T F_UU w_1 */
332: TSComputeIHessianProductFunctionUU(ts,th->stage_time,th->X,VecsDeltaLam,ts->vec_sensip_col,ts->vecs_fuu);
333: /* lambda_s^T F_UP w_2 */
334: TSComputeIHessianProductFunctionUP(ts,th->stage_time,th->X,VecsDeltaLam,ts->vec_dir,ts->vecs_fup);
335: for (nadj=0; nadj<ts->numcost; nadj++) { /* compute the residual */
336: VecCopy(ts->vecs_sensi2[nadj],VecsSensi2Temp[nadj]);
337: VecScale(VecsSensi2Temp[nadj],1./adjoint_time_step);
338: VecAXPY(VecsSensi2Temp[nadj],-1.,ts->vecs_fuu[nadj]);
339: if (ts->vecs_fup) {
340: VecAXPY(VecsSensi2Temp[nadj],-1.,ts->vecs_fup[nadj]);
341: }
342: }
343: /* Solve stage equation LHS X = RHS for second-order adjoint */
344: for (nadj=0; nadj<ts->numcost; nadj++) {
345: KSPConvergedReason kspreason;
346: KSPSolveTranspose(ksp,VecsSensi2Temp[nadj],VecsDeltaLam2[nadj]);
347: KSPGetConvergedReason(ksp,&kspreason);
348: if (kspreason < 0) {
349: ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
350: PetscInfo2(ts,"Step=%D, %Dth cost function, transposed linear solve fails, stopping 2nd-order adjoint solve\n",ts->steps,nadj);
351: }
352: }
353: }
355: /* Update sensitivities, and evaluate integrals if there is any */
356: if (!isexplicitode) {
357: th->shift = 0.0;
358: TSComputeSNESJacobian(ts,th->X,J,Jpre);
359: KSPSetOperators(ksp,J,Jpre);
360: MatScale(J,-1.);
361: for (nadj=0; nadj<ts->numcost; nadj++) {
362: /* Add f_U \lambda_s to the original RHS */
363: MatMultTransposeAdd(J,VecsDeltaLam[nadj],VecsSensiTemp[nadj],VecsSensiTemp[nadj]);
364: VecScale(VecsSensiTemp[nadj],adjoint_time_step);
365: VecCopy(VecsSensiTemp[nadj],ts->vecs_sensi[nadj]);
366: if (ts->vecs_sensi2) {
367: MatMultTransposeAdd(J,VecsDeltaLam2[nadj],VecsSensi2Temp[nadj],VecsSensi2Temp[nadj]);
368: VecScale(VecsSensi2Temp[nadj],adjoint_time_step);
369: VecCopy(VecsSensi2Temp[nadj],ts->vecs_sensi2[nadj]);
370: }
371: }
372: }
373: if (ts->vecs_sensip) {
374: TSComputeIJacobianP(ts,th->stage_time,th->X,th->Xdot,1./adjoint_time_step,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],-adjoint_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],adjoint_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],-adjoint_time_step,VecsDeltaMu2[nadj]);
398: if (ts->vecs_fpu) {
399: VecAXPY(ts->vecs_sensi2p[nadj],-adjoint_time_step,ts->vecs_fpu[nadj]);
400: }
401: if (ts->vecs_fpp) {
402: VecAXPY(ts->vecs_sensi2p[nadj],-adjoint_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;
426: PetscReal adjoint_time_step;
427: PetscReal adjoint_ptime; /* end time of the adjoint time step (ts->ptime is the start time, ususally ts->ptime is larger than adjoint_ptime) */
431: if (th->Theta == 1.) {
432: TSAdjointStepBEuler_Private(ts);
433: return(0);
434: }
435: th->status = TS_STEP_INCOMPLETE;
436: SNESGetKSP(ts->snes,&ksp);
437: TSGetIJacobian(ts,&J,&Jpre,NULL,NULL);
438: if (quadts) {
439: TSGetRHSJacobian(quadts,&quadJ,NULL,NULL,NULL);
440: TSGetRHSJacobianP(quadts,&quadJp,NULL,NULL);
441: }
442: /* If endpoint=1, th->ptime and th->X0 will be used; if endpoint=0, th->stage_time and th->X will be used. */
443: th->stage_time = th->endpoint ? ts->ptime : (ts->ptime+(1.-th->Theta)*ts->time_step);
444: adjoint_ptime = ts->ptime + ts->time_step;
445: adjoint_time_step = -ts->time_step; /* always positive since time_step is negative */
447: /* Build RHS for first-order adjoint */
448: /* Cost function has an integral term */
449: if (quadts) {
450: if (th->endpoint) {
451: TSComputeRHSJacobian(quadts,th->stage_time,ts->vec_sol,quadJ,NULL);
452: } else {
453: TSComputeRHSJacobian(quadts,th->stage_time,th->X,quadJ,NULL);
454: }
455: }
457: for (nadj=0; nadj<ts->numcost; nadj++) {
458: VecCopy(ts->vecs_sensi[nadj],VecsSensiTemp[nadj]);
459: VecScale(VecsSensiTemp[nadj],1./(th->Theta*adjoint_time_step));
460: if (quadJ) {
461: MatDenseGetColumn(quadJ,nadj,&xarr);
462: VecPlaceArray(ts->vec_drdu_col,xarr);
463: VecAXPY(VecsSensiTemp[nadj],1.,ts->vec_drdu_col);
464: VecResetArray(ts->vec_drdu_col);
465: MatDenseRestoreColumn(quadJ,&xarr);
466: }
467: }
469: /* Build LHS for first-order adjoint */
470: th->shift = 1./(th->Theta*adjoint_time_step);
471: if (th->endpoint) {
472: TSComputeSNESJacobian(ts,ts->vec_sol,J,Jpre);
473: } else {
474: TSComputeSNESJacobian(ts,th->X,J,Jpre);
475: }
476: KSPSetOperators(ksp,J,Jpre);
478: /* Solve stage equation LHS*lambda_s = RHS for first-order adjoint */
479: for (nadj=0; nadj<ts->numcost; nadj++) {
480: KSPConvergedReason kspreason;
481: KSPSolveTranspose(ksp,VecsSensiTemp[nadj],VecsDeltaLam[nadj]);
482: KSPGetConvergedReason(ksp,&kspreason);
483: if (kspreason < 0) {
484: ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
485: PetscInfo2(ts,"Step=%D, %Dth cost function, transposed linear solve fails, stopping 1st-order adjoint solve\n",ts->steps,nadj);
486: }
487: }
489: /* Second-order adjoint */
490: if (ts->vecs_sensi2) { /* U_{n+1} */
491: if (!th->endpoint) SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_SUP,"Operation not implemented in TS_Theta");
492: /* Get w1 at t_{n+1} from TLM matrix */
493: MatDenseGetColumn(ts->mat_sensip,0,&xarr);
494: VecPlaceArray(ts->vec_sensip_col,xarr);
495: /* lambda_s^T F_UU w_1 */
496: TSComputeIHessianProductFunctionUU(ts,th->stage_time,ts->vec_sol,VecsDeltaLam,ts->vec_sensip_col,ts->vecs_fuu);
497: VecResetArray(ts->vec_sensip_col);
498: MatDenseRestoreColumn(ts->mat_sensip,&xarr);
499: /* lambda_s^T F_UP w_2 */
500: TSComputeIHessianProductFunctionUP(ts,th->stage_time,ts->vec_sol,VecsDeltaLam,ts->vec_dir,ts->vecs_fup);
501: for (nadj=0; nadj<ts->numcost; nadj++) { /* compute the residual */
502: VecCopy(ts->vecs_sensi2[nadj],VecsSensi2Temp[nadj]);
503: VecScale(VecsSensi2Temp[nadj],th->shift);
504: VecAXPY(VecsSensi2Temp[nadj],-1.,ts->vecs_fuu[nadj]);
505: if (ts->vecs_fup) {
506: VecAXPY(VecsSensi2Temp[nadj],-1.,ts->vecs_fup[nadj]);
507: }
508: }
509: /* Solve stage equation LHS X = RHS for second-order adjoint */
510: for (nadj=0; nadj<ts->numcost; nadj++) {
511: KSPConvergedReason kspreason;
512: KSPSolveTranspose(ksp,VecsSensi2Temp[nadj],VecsDeltaLam2[nadj]);
513: KSPGetConvergedReason(ksp,&kspreason);
514: if (kspreason < 0) {
515: ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
516: PetscInfo2(ts,"Step=%D, %Dth cost function, transposed linear solve fails, stopping 2nd-order adjoint solve\n",ts->steps,nadj);
517: }
518: }
519: }
521: /* Update sensitivities, and evaluate integrals if there is any */
522: if(th->endpoint) { /* two-stage Theta methods with th->Theta!=1, th->Theta==1 leads to BEuler */
523: th->shift = 1./((th->Theta-1.)*adjoint_time_step);
524: th->stage_time = adjoint_ptime;
525: TSComputeSNESJacobian(ts,th->X0,J,Jpre);
526: KSPSetOperators(ksp,J,Jpre);
527: /* R_U at t_n */
528: if (quadts) {
529: TSComputeRHSJacobian(quadts,adjoint_ptime,th->X0,quadJ,NULL);
530: }
531: for (nadj=0; nadj<ts->numcost; nadj++) {
532: MatMultTranspose(J,VecsDeltaLam[nadj],ts->vecs_sensi[nadj]);
533: if (quadJ) {
534: MatDenseGetColumn(quadJ,nadj,&xarr);
535: VecPlaceArray(ts->vec_drdu_col,xarr);
536: VecAXPY(ts->vecs_sensi[nadj],-1.,ts->vec_drdu_col);
537: VecResetArray(ts->vec_drdu_col);
538: MatDenseRestoreColumn(quadJ,&xarr);
539: }
540: VecScale(ts->vecs_sensi[nadj],1./th->shift);
541: }
543: /* Second-order adjoint */
544: if (ts->vecs_sensi2) { /* U_n */
545: /* Get w1 at t_n from TLM matrix */
546: MatDenseGetColumn(th->MatFwdSensip0,0,&xarr);
547: VecPlaceArray(ts->vec_sensip_col,xarr);
548: /* lambda_s^T F_UU w_1 */
549: TSComputeIHessianProductFunctionUU(ts,adjoint_ptime,th->X0,VecsDeltaLam,ts->vec_sensip_col,ts->vecs_fuu);
550: VecResetArray(ts->vec_sensip_col);
551: MatDenseRestoreColumn(th->MatFwdSensip0,&xarr);
552: /* lambda_s^T F_UU w_2 */
553: TSComputeIHessianProductFunctionUP(ts,adjoint_ptime,th->X0,VecsDeltaLam,ts->vec_dir,ts->vecs_fup);
554: for (nadj=0; nadj<ts->numcost; nadj++) {
555: /* 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 */
556: MatMultTranspose(J,VecsDeltaLam2[nadj],ts->vecs_sensi2[nadj]);
557: VecAXPY(ts->vecs_sensi2[nadj],1.,ts->vecs_fuu[nadj]);
558: if (ts->vecs_fup) {
559: VecAXPY(ts->vecs_sensi2[nadj],1.,ts->vecs_fup[nadj]);
560: }
561: VecScale(ts->vecs_sensi2[nadj],1./th->shift);
562: }
563: }
565: th->stage_time = ts->ptime; /* recover the old value */
567: if (ts->vecs_sensip) { /* sensitivities wrt parameters */
568: /* U_{n+1} */
569: TSComputeIJacobianP(ts,th->stage_time,ts->vec_sol,th->Xdot,-1./(th->Theta*adjoint_time_step),ts->Jacp,PETSC_FALSE);
570: if (quadts) {
571: TSComputeRHSJacobianP(quadts,th->stage_time,ts->vec_sol,quadJp);
572: }
573: for (nadj=0; nadj<ts->numcost; nadj++) {
574: MatMultTranspose(ts->Jacp,VecsDeltaLam[nadj],VecsDeltaMu[nadj]);
575: VecAXPY(ts->vecs_sensip[nadj],-adjoint_time_step*th->Theta,VecsDeltaMu[nadj]);
576: }
577: if (ts->vecs_sensi2p) { /* second-order */
578: /* Get w1 at t_{n+1} from TLM matrix */
579: MatDenseGetColumn(ts->mat_sensip,0,&xarr);
580: VecPlaceArray(ts->vec_sensip_col,xarr);
581: /* lambda_s^T F_PU w_1 */
582: TSComputeIHessianProductFunctionPU(ts,th->stage_time,ts->vec_sol,VecsDeltaLam,ts->vec_sensip_col,ts->vecs_fpu);
583: VecResetArray(ts->vec_sensip_col);
584: MatDenseRestoreColumn(ts->mat_sensip,&xarr);
586: /* lambda_s^T F_PP w_2 */
587: TSComputeIHessianProductFunctionPP(ts,th->stage_time,ts->vec_sol,VecsDeltaLam,ts->vec_dir,ts->vecs_fpp);
588: for (nadj=0; nadj<ts->numcost; nadj++) {
589: /* 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) */
590: MatMultTranspose(ts->Jacp,VecsDeltaLam2[nadj],VecsDeltaMu2[nadj]);
591: VecAXPY(ts->vecs_sensi2p[nadj],-adjoint_time_step*th->Theta,VecsDeltaMu2[nadj]);
592: if (ts->vecs_fpu) {
593: VecAXPY(ts->vecs_sensi2p[nadj],-adjoint_time_step*th->Theta,ts->vecs_fpu[nadj]);
594: }
595: if (ts->vecs_fpp) {
596: VecAXPY(ts->vecs_sensi2p[nadj],-adjoint_time_step*th->Theta,ts->vecs_fpp[nadj]);
597: }
598: }
599: }
601: /* U_s */
602: TSComputeIJacobianP(ts,adjoint_ptime,th->X0,th->Xdot,1./((th->Theta-1.0)*adjoint_time_step),ts->Jacp,PETSC_FALSE);
603: if (quadts) {
604: TSComputeRHSJacobianP(quadts,adjoint_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],-adjoint_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,adjoint_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,adjoint_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],-adjoint_time_step*(1.0-th->Theta),VecsDeltaMu2[nadj]);
623: if (ts->vecs_fpu) {
624: VecAXPY(ts->vecs_sensi2p[nadj],-adjoint_time_step*(1.0-th->Theta),ts->vecs_fpu[nadj]);
625: }
626: if (ts->vecs_fpp) {
627: VecAXPY(ts->vecs_sensi2p[nadj],-adjoint_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],-adjoint_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],adjoint_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],-adjoint_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],adjoint_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;
745: PetscReal previous_shift;
749: previous_shift = th->shift;
750: MatCopy(ts->mat_sensip,th->MatFwdSensip0,SAME_NONZERO_PATTERN);
752: if (quadts && quadts->mat_sensip) {
753: MatCopy(quadts->mat_sensip,th->MatIntegralSensip0,SAME_NONZERO_PATTERN);
754: }
755: SNESGetKSP(ts->snes,&ksp);
756: TSGetIJacobian(ts,&J,&Jpre,NULL,NULL);
757: if (quadts) {
758: TSGetRHSJacobian(quadts,&quadJ,NULL,NULL,NULL);
759: TSGetRHSJacobianP(quadts,&quadJp,NULL,NULL);
760: }
762: /* Build RHS */
763: if (th->endpoint) { /* 2-stage method*/
764: th->shift = 1./((th->Theta-1.)*th->time_step0);
765: TSComputeIJacobian(ts,th->ptime0,th->X0,th->Xdot,th->shift,J,Jpre,PETSC_FALSE);
766: MatMatMult(J,ts->mat_sensip,MAT_REUSE_MATRIX,PETSC_DEFAULT,&MatDeltaFwdSensip);
767: MatScale(MatDeltaFwdSensip,(th->Theta-1.)/th->Theta);
769: /* Add the f_p forcing terms */
770: if (ts->Jacp) {
771: TSComputeIJacobianP(ts,th->ptime0,th->X0,th->Xdot,th->shift,ts->Jacp,PETSC_FALSE);
772: MatAXPY(MatDeltaFwdSensip,(th->Theta-1.)/th->Theta,ts->Jacp,SUBSET_NONZERO_PATTERN);
773: TSComputeIJacobianP(ts,th->stage_time,ts->vec_sol,th->Xdot,th->shift,ts->Jacp,PETSC_FALSE);
774: MatAXPY(MatDeltaFwdSensip,-1.,ts->Jacp,SUBSET_NONZERO_PATTERN);
775: }
776: } else { /* 1-stage method */
777: th->shift = 0.0;
778: TSComputeIJacobian(ts,th->stage_time,th->X,th->Xdot,th->shift,J,Jpre,PETSC_FALSE);
779: MatMatMult(J,ts->mat_sensip,MAT_REUSE_MATRIX,PETSC_DEFAULT,&MatDeltaFwdSensip);
780: MatScale(MatDeltaFwdSensip,-1.);
782: /* Add the f_p forcing terms */
783: if (ts->Jacp) {
784: TSComputeIJacobianP(ts,th->stage_time,th->X,th->Xdot,th->shift,ts->Jacp,PETSC_FALSE);
785: MatAXPY(MatDeltaFwdSensip,-1.,ts->Jacp,SUBSET_NONZERO_PATTERN);
786: }
787: }
789: /* Build LHS */
790: th->shift = previous_shift; /* recover the previous shift used in TSStep_Theta() */
791: if (th->endpoint) {
792: TSComputeIJacobian(ts,th->stage_time,ts->vec_sol,th->Xdot,th->shift,J,Jpre,PETSC_FALSE);
793: } else {
794: TSComputeIJacobian(ts,th->stage_time,th->X,th->Xdot,th->shift,J,Jpre,PETSC_FALSE);
795: }
796: KSPSetOperators(ksp,J,Jpre);
798: /*
799: Evaluate the first stage of integral gradients with the 2-stage method:
800: drdu|t_n*S(t_n) + drdp|t_n
801: This is done before the linear solve because the sensitivity variable S(t_n) will be propagated to S(t_{n+1})
802: */
803: if (th->endpoint) { /* 2-stage method only */
804: if (quadts && quadts->mat_sensip) {
805: TSComputeRHSJacobian(quadts,th->ptime0,th->X0,quadJ,NULL);
806: TSComputeRHSJacobianP(quadts,th->ptime0,th->X0,quadJp);
807: MatTransposeMatMult(ts->mat_sensip,quadJ,MAT_REUSE_MATRIX,PETSC_DEFAULT,&th->MatIntegralSensipTemp);
808: MatAXPY(th->MatIntegralSensipTemp,1,quadJp,SAME_NONZERO_PATTERN);
809: MatAXPY(quadts->mat_sensip,th->time_step0*(1.-th->Theta),th->MatIntegralSensipTemp,SAME_NONZERO_PATTERN);
810: }
811: }
813: /* Solve the tangent linear equation for forward sensitivities to parameters */
814: for (ntlm=0; ntlm<th->num_tlm; ntlm++) {
815: KSPConvergedReason kspreason;
816: MatDenseGetColumn(MatDeltaFwdSensip,ntlm,&barr);
817: VecPlaceArray(VecDeltaFwdSensipCol,barr);
818: if (th->endpoint) {
819: MatDenseGetColumn(ts->mat_sensip,ntlm,&xarr);
820: VecPlaceArray(ts->vec_sensip_col,xarr);
821: KSPSolve(ksp,VecDeltaFwdSensipCol,ts->vec_sensip_col);
822: VecResetArray(ts->vec_sensip_col);
823: MatDenseRestoreColumn(ts->mat_sensip,&xarr);
824: } else {
825: KSPSolve(ksp,VecDeltaFwdSensipCol,VecDeltaFwdSensipCol);
826: }
827: KSPGetConvergedReason(ksp,&kspreason);
828: if (kspreason < 0) {
829: ts->reason = TSFORWARD_DIVERGED_LINEAR_SOLVE;
830: PetscInfo2(ts,"Step=%D, %Dth tangent linear solve, linear solve fails, stopping tangent linear solve\n",ts->steps,ntlm);
831: }
832: VecResetArray(VecDeltaFwdSensipCol);
833: MatDenseRestoreColumn(MatDeltaFwdSensip,&barr);
834: }
836: /*
837: Evaluate the second stage of integral gradients with the 2-stage method:
838: drdu|t_{n+1}*S(t_{n+1}) + drdp|t_{n+1}
839: */
840: if (quadts && quadts->mat_sensip) {
841: if (!th->endpoint) {
842: MatAXPY(ts->mat_sensip,1,MatDeltaFwdSensip,SAME_NONZERO_PATTERN); /* stage sensitivity */
843: TSComputeRHSJacobian(quadts,th->stage_time,th->X,quadJ,NULL);
844: TSComputeRHSJacobianP(quadts,th->stage_time,th->X,quadJp);
845: MatTransposeMatMult(ts->mat_sensip,quadJ,MAT_REUSE_MATRIX,PETSC_DEFAULT,&th->MatIntegralSensipTemp);
846: MatAXPY(th->MatIntegralSensipTemp,1,quadJp,SAME_NONZERO_PATTERN);
847: MatAXPY(quadts->mat_sensip,th->time_step0,th->MatIntegralSensipTemp,SAME_NONZERO_PATTERN);
848: MatAXPY(ts->mat_sensip,(1.-th->Theta)/th->Theta,MatDeltaFwdSensip,SAME_NONZERO_PATTERN);
849: } else {
850: TSComputeRHSJacobian(quadts,th->stage_time,ts->vec_sol,quadJ,NULL);
851: TSComputeRHSJacobianP(quadts,th->stage_time,ts->vec_sol,quadJp);
852: MatTransposeMatMult(ts->mat_sensip,quadJ,MAT_REUSE_MATRIX,PETSC_DEFAULT,&th->MatIntegralSensipTemp);
853: MatAXPY(th->MatIntegralSensipTemp,1,quadJp,SAME_NONZERO_PATTERN);
854: MatAXPY(quadts->mat_sensip,th->time_step0*th->Theta,th->MatIntegralSensipTemp,SAME_NONZERO_PATTERN);
855: }
856: } else {
857: if (!th->endpoint) {
858: MatAXPY(ts->mat_sensip,1./th->Theta,MatDeltaFwdSensip,SAME_NONZERO_PATTERN);
859: }
860: }
861: return(0);
862: }
864: static PetscErrorCode TSForwardGetStages_Theta(TS ts,PetscInt *ns,Mat **stagesensip)
865: {
866: TS_Theta *th = (TS_Theta*)ts->data;
869: if (ns) *ns = 1;
870: if (stagesensip) *stagesensip = th->endpoint ? &(th->MatFwdSensip0) : &(th->MatDeltaFwdSensip);
871: return(0);
872: }
874: /*------------------------------------------------------------*/
875: static PetscErrorCode TSReset_Theta(TS ts)
876: {
877: TS_Theta *th = (TS_Theta*)ts->data;
881: VecDestroy(&th->X);
882: VecDestroy(&th->Xdot);
883: VecDestroy(&th->X0);
884: VecDestroy(&th->affine);
886: VecDestroy(&th->vec_sol_prev);
887: VecDestroy(&th->vec_lte_work);
889: VecDestroy(&th->VecCostIntegral0);
890: return(0);
891: }
893: static PetscErrorCode TSAdjointReset_Theta(TS ts)
894: {
895: TS_Theta *th = (TS_Theta*)ts->data;
899: VecDestroyVecs(ts->numcost,&th->VecsDeltaLam);
900: VecDestroyVecs(ts->numcost,&th->VecsDeltaMu);
901: VecDestroyVecs(ts->numcost,&th->VecsDeltaLam2);
902: VecDestroyVecs(ts->numcost,&th->VecsDeltaMu2);
903: VecDestroyVecs(ts->numcost,&th->VecsSensiTemp);
904: VecDestroyVecs(ts->numcost,&th->VecsSensi2Temp);
905: return(0);
906: }
908: static PetscErrorCode TSDestroy_Theta(TS ts)
909: {
913: TSReset_Theta(ts);
914: if (ts->dm) {
915: DMCoarsenHookRemove(ts->dm,DMCoarsenHook_TSTheta,DMRestrictHook_TSTheta,ts);
916: DMSubDomainHookRemove(ts->dm,DMSubDomainHook_TSTheta,DMSubDomainRestrictHook_TSTheta,ts);
917: }
918: PetscFree(ts->data);
919: PetscObjectComposeFunction((PetscObject)ts,"TSThetaGetTheta_C",NULL);
920: PetscObjectComposeFunction((PetscObject)ts,"TSThetaSetTheta_C",NULL);
921: PetscObjectComposeFunction((PetscObject)ts,"TSThetaGetEndpoint_C",NULL);
922: PetscObjectComposeFunction((PetscObject)ts,"TSThetaSetEndpoint_C",NULL);
923: return(0);
924: }
926: /*
927: This defines the nonlinear equation that is to be solved with SNES
928: G(U) = F[t0+Theta*dt, U, (U-U0)*shift] = 0
929: */
930: static PetscErrorCode SNESTSFormFunction_Theta(SNES snes,Vec x,Vec y,TS ts)
931: {
932: TS_Theta *th = (TS_Theta*)ts->data;
934: Vec X0,Xdot;
935: DM dm,dmsave;
936: PetscReal shift = th->shift;
939: SNESGetDM(snes,&dm);
940: /* When using the endpoint variant, this is actually 1/Theta * Xdot */
941: TSThetaGetX0AndXdot(ts,dm,&X0,&Xdot);
942: if (x != X0) {
943: VecAXPBYPCZ(Xdot,-shift,shift,0,X0,x);
944: } else {
945: VecZeroEntries(Xdot);
946: }
947: /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */
948: dmsave = ts->dm;
949: ts->dm = dm;
950: TSComputeIFunction(ts,th->stage_time,x,Xdot,y,PETSC_FALSE);
951: ts->dm = dmsave;
952: TSThetaRestoreX0AndXdot(ts,dm,&X0,&Xdot);
953: return(0);
954: }
956: static PetscErrorCode SNESTSFormJacobian_Theta(SNES snes,Vec x,Mat A,Mat B,TS ts)
957: {
958: TS_Theta *th = (TS_Theta*)ts->data;
960: Vec Xdot;
961: DM dm,dmsave;
962: PetscReal shift = th->shift;
965: SNESGetDM(snes,&dm);
966: /* Xdot has already been computed in SNESTSFormFunction_Theta (SNES guarantees this) */
967: TSThetaGetX0AndXdot(ts,dm,NULL,&Xdot);
969: dmsave = ts->dm;
970: ts->dm = dm;
971: TSComputeIJacobian(ts,th->stage_time,x,Xdot,shift,A,B,PETSC_FALSE);
972: ts->dm = dmsave;
973: TSThetaRestoreX0AndXdot(ts,dm,NULL,&Xdot);
974: return(0);
975: }
977: static PetscErrorCode TSForwardSetUp_Theta(TS ts)
978: {
979: TS_Theta *th = (TS_Theta*)ts->data;
980: TS quadts = ts->quadraturets;
984: /* combine sensitivities to parameters and sensitivities to initial values into one array */
985: th->num_tlm = ts->num_parameters;
986: MatDuplicate(ts->mat_sensip,MAT_DO_NOT_COPY_VALUES,&th->MatDeltaFwdSensip);
987: if (quadts && quadts->mat_sensip) {
988: MatDuplicate(quadts->mat_sensip,MAT_DO_NOT_COPY_VALUES,&th->MatIntegralSensipTemp);
989: MatDuplicate(quadts->mat_sensip,MAT_DO_NOT_COPY_VALUES,&th->MatIntegralSensip0);
990: }
991: /* backup sensitivity results for roll-backs */
992: MatDuplicate(ts->mat_sensip,MAT_DO_NOT_COPY_VALUES,&th->MatFwdSensip0);
994: VecDuplicate(ts->vec_sol,&th->VecDeltaFwdSensipCol);
995: return(0);
996: }
998: static PetscErrorCode TSForwardReset_Theta(TS ts)
999: {
1000: TS_Theta *th = (TS_Theta*)ts->data;
1001: TS quadts = ts->quadraturets;
1005: if (quadts && quadts->mat_sensip) {
1006: MatDestroy(&th->MatIntegralSensipTemp);
1007: MatDestroy(&th->MatIntegralSensip0);
1008: }
1009: VecDestroy(&th->VecDeltaFwdSensipCol);
1010: MatDestroy(&th->MatDeltaFwdSensip);
1011: MatDestroy(&th->MatFwdSensip0);
1012: return(0);
1013: }
1015: static PetscErrorCode TSSetUp_Theta(TS ts)
1016: {
1017: TS_Theta *th = (TS_Theta*)ts->data;
1018: TS quadts = ts->quadraturets;
1019: PetscBool match;
1023: if (!th->VecCostIntegral0 && quadts && ts->costintegralfwd) { /* back up cost integral */
1024: VecDuplicate(quadts->vec_sol,&th->VecCostIntegral0);
1025: }
1026: if (!th->X) {
1027: VecDuplicate(ts->vec_sol,&th->X);
1028: }
1029: if (!th->Xdot) {
1030: VecDuplicate(ts->vec_sol,&th->Xdot);
1031: }
1032: if (!th->X0) {
1033: VecDuplicate(ts->vec_sol,&th->X0);
1034: }
1035: if (th->endpoint) {
1036: VecDuplicate(ts->vec_sol,&th->affine);
1037: }
1039: th->order = (th->Theta == 0.5) ? 2 : 1;
1040: th->shift = 1/(th->Theta*ts->time_step);
1042: TSGetDM(ts,&ts->dm);
1043: DMCoarsenHookAdd(ts->dm,DMCoarsenHook_TSTheta,DMRestrictHook_TSTheta,ts);
1044: DMSubDomainHookAdd(ts->dm,DMSubDomainHook_TSTheta,DMSubDomainRestrictHook_TSTheta,ts);
1046: TSGetAdapt(ts,&ts->adapt);
1047: TSAdaptCandidatesClear(ts->adapt);
1048: PetscObjectTypeCompare((PetscObject)ts->adapt,TSADAPTNONE,&match);
1049: if (!match) {
1050: VecDuplicate(ts->vec_sol,&th->vec_sol_prev);
1051: VecDuplicate(ts->vec_sol,&th->vec_lte_work);
1052: }
1053: TSGetSNES(ts,&ts->snes);
1054: return(0);
1055: }
1057: /*------------------------------------------------------------*/
1059: static PetscErrorCode TSAdjointSetUp_Theta(TS ts)
1060: {
1061: TS_Theta *th = (TS_Theta*)ts->data;
1065: VecDuplicateVecs(ts->vecs_sensi[0],ts->numcost,&th->VecsDeltaLam);
1066: VecDuplicateVecs(ts->vecs_sensi[0],ts->numcost,&th->VecsSensiTemp);
1067: if (ts->vecs_sensip) {
1068: VecDuplicateVecs(ts->vecs_sensip[0],ts->numcost,&th->VecsDeltaMu);
1069: }
1070: if (ts->vecs_sensi2) {
1071: VecDuplicateVecs(ts->vecs_sensi[0],ts->numcost,&th->VecsDeltaLam2);
1072: VecDuplicateVecs(ts->vecs_sensi2[0],ts->numcost,&th->VecsSensi2Temp);
1073: /* hack ts to make implicit TS solver work when users provide only explicit versions of callbacks (RHSFunction,RHSJacobian,RHSHessian etc.) */
1074: if (!ts->ihessianproduct_fuu) ts->vecs_fuu = ts->vecs_guu;
1075: if (!ts->ihessianproduct_fup) ts->vecs_fup = ts->vecs_gup;
1076: }
1077: if (ts->vecs_sensi2p) {
1078: VecDuplicateVecs(ts->vecs_sensi2p[0],ts->numcost,&th->VecsDeltaMu2);
1079: /* hack ts to make implicit TS solver work when users provide only explicit versions of callbacks (RHSFunction,RHSJacobian,RHSHessian etc.) */
1080: if (!ts->ihessianproduct_fpu) ts->vecs_fpu = ts->vecs_gpu;
1081: if (!ts->ihessianproduct_fpp) ts->vecs_fpp = ts->vecs_gpp;
1082: }
1083: return(0);
1084: }
1086: static PetscErrorCode TSSetFromOptions_Theta(PetscOptionItems *PetscOptionsObject,TS ts)
1087: {
1088: TS_Theta *th = (TS_Theta*)ts->data;
1092: PetscOptionsHead(PetscOptionsObject,"Theta ODE solver options");
1093: {
1094: PetscOptionsReal("-ts_theta_theta","Location of stage (0<Theta<=1)","TSThetaSetTheta",th->Theta,&th->Theta,NULL);
1095: PetscOptionsBool("-ts_theta_endpoint","Use the endpoint instead of midpoint form of the Theta method","TSThetaSetEndpoint",th->endpoint,&th->endpoint,NULL);
1096: PetscOptionsBool("-ts_theta_initial_guess_extrapolate","Extrapolate stage initial guess from previous solution (sometimes unstable)","TSThetaSetExtrapolate",th->extrapolate,&th->extrapolate,NULL);
1097: }
1098: PetscOptionsTail();
1099: return(0);
1100: }
1102: static PetscErrorCode TSView_Theta(TS ts,PetscViewer viewer)
1103: {
1104: TS_Theta *th = (TS_Theta*)ts->data;
1105: PetscBool iascii;
1109: PetscObjectTypeCompare((PetscObject)viewer,PETSCVIEWERASCII,&iascii);
1110: if (iascii) {
1111: PetscViewerASCIIPrintf(viewer," Theta=%g\n",(double)th->Theta);
1112: PetscViewerASCIIPrintf(viewer," Extrapolation=%s\n",th->extrapolate ? "yes" : "no");
1113: }
1114: return(0);
1115: }
1117: static PetscErrorCode TSThetaGetTheta_Theta(TS ts,PetscReal *theta)
1118: {
1119: TS_Theta *th = (TS_Theta*)ts->data;
1122: *theta = th->Theta;
1123: return(0);
1124: }
1126: static PetscErrorCode TSThetaSetTheta_Theta(TS ts,PetscReal theta)
1127: {
1128: TS_Theta *th = (TS_Theta*)ts->data;
1131: if (theta <= 0 || 1 < theta) SETERRQ1(PetscObjectComm((PetscObject)ts),PETSC_ERR_ARG_OUTOFRANGE,"Theta %g not in range (0,1]",(double)theta);
1132: th->Theta = theta;
1133: th->order = (th->Theta == 0.5) ? 2 : 1;
1134: return(0);
1135: }
1137: static PetscErrorCode TSThetaGetEndpoint_Theta(TS ts,PetscBool *endpoint)
1138: {
1139: TS_Theta *th = (TS_Theta*)ts->data;
1142: *endpoint = th->endpoint;
1143: return(0);
1144: }
1146: static PetscErrorCode TSThetaSetEndpoint_Theta(TS ts,PetscBool flg)
1147: {
1148: TS_Theta *th = (TS_Theta*)ts->data;
1151: th->endpoint = flg;
1152: return(0);
1153: }
1155: #if defined(PETSC_HAVE_COMPLEX)
1156: static PetscErrorCode TSComputeLinearStability_Theta(TS ts,PetscReal xr,PetscReal xi,PetscReal *yr,PetscReal *yi)
1157: {
1158: PetscComplex z = xr + xi*PETSC_i,f;
1159: TS_Theta *th = (TS_Theta*)ts->data;
1160: const PetscReal one = 1.0;
1163: f = (one + (one - th->Theta)*z)/(one - th->Theta*z);
1164: *yr = PetscRealPartComplex(f);
1165: *yi = PetscImaginaryPartComplex(f);
1166: return(0);
1167: }
1168: #endif
1170: static PetscErrorCode TSGetStages_Theta(TS ts,PetscInt *ns,Vec **Y)
1171: {
1172: TS_Theta *th = (TS_Theta*)ts->data;
1175: if (ns) *ns = 1;
1176: if (Y) *Y = th->endpoint ? &(th->X0) : &(th->X);
1177: return(0);
1178: }
1180: /* ------------------------------------------------------------ */
1181: /*MC
1182: TSTHETA - DAE solver using the implicit Theta method
1184: Level: beginner
1186: Options Database:
1187: + -ts_theta_theta <Theta> - Location of stage (0<Theta<=1)
1188: . -ts_theta_endpoint <flag> - Use the endpoint (like Crank-Nicholson) instead of midpoint form of the Theta method
1189: - -ts_theta_initial_guess_extrapolate <flg> - Extrapolate stage initial guess from previous solution (sometimes unstable)
1191: Notes:
1192: $ -ts_type theta -ts_theta_theta 1.0 corresponds to backward Euler (TSBEULER)
1193: $ -ts_type theta -ts_theta_theta 0.5 corresponds to the implicit midpoint rule
1194: $ -ts_type theta -ts_theta_theta 0.5 -ts_theta_endpoint corresponds to Crank-Nicholson (TSCN)
1196: This method can be applied to DAE.
1198: This method is cast as a 1-stage implicit Runge-Kutta method.
1200: .vb
1201: Theta | Theta
1202: -------------
1203: | 1
1204: .ve
1206: For the default Theta=0.5, this is also known as the implicit midpoint rule.
1208: When the endpoint variant is chosen, the method becomes a 2-stage method with first stage explicit:
1210: .vb
1211: 0 | 0 0
1212: 1 | 1-Theta Theta
1213: -------------------
1214: | 1-Theta Theta
1215: .ve
1217: For the default Theta=0.5, this is the trapezoid rule (also known as Crank-Nicolson, see TSCN).
1219: To apply a diagonally implicit RK method to DAE, the stage formula
1221: $ Y_i = X + h sum_j a_ij Y'_j
1223: is interpreted as a formula for Y'_i in terms of Y_i and known values (Y'_j, j<i)
1225: .seealso: TSCreate(), TS, TSSetType(), TSCN, TSBEULER, TSThetaSetTheta(), TSThetaSetEndpoint()
1227: M*/
1228: PETSC_EXTERN PetscErrorCode TSCreate_Theta(TS ts)
1229: {
1230: TS_Theta *th;
1234: ts->ops->reset = TSReset_Theta;
1235: ts->ops->adjointreset = TSAdjointReset_Theta;
1236: ts->ops->destroy = TSDestroy_Theta;
1237: ts->ops->view = TSView_Theta;
1238: ts->ops->setup = TSSetUp_Theta;
1239: ts->ops->adjointsetup = TSAdjointSetUp_Theta;
1240: ts->ops->adjointreset = TSAdjointReset_Theta;
1241: ts->ops->step = TSStep_Theta;
1242: ts->ops->interpolate = TSInterpolate_Theta;
1243: ts->ops->evaluatewlte = TSEvaluateWLTE_Theta;
1244: ts->ops->rollback = TSRollBack_Theta;
1245: ts->ops->setfromoptions = TSSetFromOptions_Theta;
1246: ts->ops->snesfunction = SNESTSFormFunction_Theta;
1247: ts->ops->snesjacobian = SNESTSFormJacobian_Theta;
1248: #if defined(PETSC_HAVE_COMPLEX)
1249: ts->ops->linearstability = TSComputeLinearStability_Theta;
1250: #endif
1251: ts->ops->getstages = TSGetStages_Theta;
1252: ts->ops->adjointstep = TSAdjointStep_Theta;
1253: ts->ops->adjointintegral = TSAdjointCostIntegral_Theta;
1254: ts->ops->forwardintegral = TSForwardCostIntegral_Theta;
1255: ts->default_adapt_type = TSADAPTNONE;
1257: ts->ops->forwardsetup = TSForwardSetUp_Theta;
1258: ts->ops->forwardreset = TSForwardReset_Theta;
1259: ts->ops->forwardstep = TSForwardStep_Theta;
1260: ts->ops->forwardgetstages = TSForwardGetStages_Theta;
1262: ts->usessnes = PETSC_TRUE;
1264: PetscNewLog(ts,&th);
1265: ts->data = (void*)th;
1267: th->VecsDeltaLam = NULL;
1268: th->VecsDeltaMu = NULL;
1269: th->VecsSensiTemp = NULL;
1270: th->VecsSensi2Temp = NULL;
1272: th->extrapolate = PETSC_FALSE;
1273: th->Theta = 0.5;
1274: th->order = 2;
1275: PetscObjectComposeFunction((PetscObject)ts,"TSThetaGetTheta_C",TSThetaGetTheta_Theta);
1276: PetscObjectComposeFunction((PetscObject)ts,"TSThetaSetTheta_C",TSThetaSetTheta_Theta);
1277: PetscObjectComposeFunction((PetscObject)ts,"TSThetaGetEndpoint_C",TSThetaGetEndpoint_Theta);
1278: PetscObjectComposeFunction((PetscObject)ts,"TSThetaSetEndpoint_C",TSThetaSetEndpoint_Theta);
1279: return(0);
1280: }
1282: /*@
1283: TSThetaGetTheta - Get the abscissa of the stage in (0,1].
1285: Not Collective
1287: Input Parameter:
1288: . ts - timestepping context
1290: Output Parameter:
1291: . theta - stage abscissa
1293: Note:
1294: Use of this function is normally only required to hack TSTHETA to use a modified integration scheme.
1296: Level: Advanced
1298: .seealso: TSThetaSetTheta()
1299: @*/
1300: PetscErrorCode TSThetaGetTheta(TS ts,PetscReal *theta)
1301: {
1307: PetscUseMethod(ts,"TSThetaGetTheta_C",(TS,PetscReal*),(ts,theta));
1308: return(0);
1309: }
1311: /*@
1312: TSThetaSetTheta - Set the abscissa of the stage in (0,1].
1314: Not Collective
1316: Input Parameter:
1317: + ts - timestepping context
1318: - theta - stage abscissa
1320: Options Database:
1321: . -ts_theta_theta <theta>
1323: Level: Intermediate
1325: .seealso: TSThetaGetTheta()
1326: @*/
1327: PetscErrorCode TSThetaSetTheta(TS ts,PetscReal theta)
1328: {
1333: PetscTryMethod(ts,"TSThetaSetTheta_C",(TS,PetscReal),(ts,theta));
1334: return(0);
1335: }
1337: /*@
1338: TSThetaGetEndpoint - Gets whether to use the endpoint variant of the method (e.g. trapezoid/Crank-Nicolson instead of midpoint rule).
1340: Not Collective
1342: Input Parameter:
1343: . ts - timestepping context
1345: Output Parameter:
1346: . endpoint - PETSC_TRUE when using the endpoint variant
1348: Level: Advanced
1350: .seealso: TSThetaSetEndpoint(), TSTHETA, TSCN
1351: @*/
1352: PetscErrorCode TSThetaGetEndpoint(TS ts,PetscBool *endpoint)
1353: {
1359: PetscUseMethod(ts,"TSThetaGetEndpoint_C",(TS,PetscBool*),(ts,endpoint));
1360: return(0);
1361: }
1363: /*@
1364: TSThetaSetEndpoint - Sets whether to use the endpoint variant of the method (e.g. trapezoid/Crank-Nicolson instead of midpoint rule).
1366: Not Collective
1368: Input Parameter:
1369: + ts - timestepping context
1370: - flg - PETSC_TRUE to use the endpoint variant
1372: Options Database:
1373: . -ts_theta_endpoint <flg>
1375: Level: Intermediate
1377: .seealso: TSTHETA, TSCN
1378: @*/
1379: PetscErrorCode TSThetaSetEndpoint(TS ts,PetscBool flg)
1380: {
1385: PetscTryMethod(ts,"TSThetaSetEndpoint_C",(TS,PetscBool),(ts,flg));
1386: return(0);
1387: }
1389: /*
1390: * TSBEULER and TSCN are straightforward specializations of TSTHETA.
1391: * The creation functions for these specializations are below.
1392: */
1394: static PetscErrorCode TSSetUp_BEuler(TS ts)
1395: {
1396: TS_Theta *th = (TS_Theta*)ts->data;
1400: 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");
1401: 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");
1402: TSSetUp_Theta(ts);
1403: return(0);
1404: }
1406: static PetscErrorCode TSView_BEuler(TS ts,PetscViewer viewer)
1407: {
1409: return(0);
1410: }
1412: /*MC
1413: TSBEULER - ODE solver using the implicit backward Euler method
1415: Level: beginner
1417: Notes:
1418: TSBEULER is equivalent to TSTHETA with Theta=1.0
1420: $ -ts_type theta -ts_theta_theta 1.0
1422: .seealso: TSCreate(), TS, TSSetType(), TSEULER, TSCN, TSTHETA
1424: M*/
1425: PETSC_EXTERN PetscErrorCode TSCreate_BEuler(TS ts)
1426: {
1430: TSCreate_Theta(ts);
1431: TSThetaSetTheta(ts,1.0);
1432: TSThetaSetEndpoint(ts,PETSC_FALSE);
1433: ts->ops->setup = TSSetUp_BEuler;
1434: ts->ops->view = TSView_BEuler;
1435: return(0);
1436: }
1438: static PetscErrorCode TSSetUp_CN(TS ts)
1439: {
1440: TS_Theta *th = (TS_Theta*)ts->data;
1444: 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");
1445: 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");
1446: TSSetUp_Theta(ts);
1447: return(0);
1448: }
1450: static PetscErrorCode TSView_CN(TS ts,PetscViewer viewer)
1451: {
1453: return(0);
1454: }
1456: /*MC
1457: TSCN - ODE solver using the implicit Crank-Nicolson method.
1459: Level: beginner
1461: Notes:
1462: TSCN is equivalent to TSTHETA with Theta=0.5 and the "endpoint" option set. I.e.
1464: $ -ts_type theta -ts_theta_theta 0.5 -ts_theta_endpoint
1466: .seealso: TSCreate(), TS, TSSetType(), TSBEULER, TSTHETA
1468: M*/
1469: PETSC_EXTERN PetscErrorCode TSCreate_CN(TS ts)
1470: {
1474: TSCreate_Theta(ts);
1475: TSThetaSetTheta(ts,0.5);
1476: TSThetaSetEndpoint(ts,PETSC_TRUE);
1477: ts->ops->setup = TSSetUp_CN;
1478: ts->ops->view = TSView_CN;
1479: return(0);
1480: }