Actual source code: theta.c

      1: /*
      2:   Code for timestepping with implicit Theta method
      3: */
      4: #include <petsc/private/tsimpl.h>
      5: #include <petscsnes.h>
      6: #include <petscdm.h>
      7: #include <petscmat.h>

      9: typedef struct {
     10:   /* context for time stepping */
     11:   PetscReal    stage_time;
     12:   Vec          Stages[2];   /* Storage for stage solutions */
     13:   Vec          X0, X, Xdot; /* Storage for u^n, u^n + dt a_{11} k_1, and time derivative u^{n+1}_t */
     14:   Vec          affine;      /* Affine vector needed for residual at beginning of step in endpoint formulation */
     15:   PetscReal    Theta;
     16:   PetscReal    shift; /* Shift parameter for SNES Jacobian, used by forward, TLM and adjoint */
     17:   PetscInt     order;
     18:   PetscBool    endpoint;
     19:   PetscBool    extrapolate;
     20:   TSStepStatus status;
     21:   Vec          VecCostIntegral0; /* Backup for roll-backs due to events, used by cost integral */
     22:   PetscReal    ptime0;           /* Backup for ts->ptime, the start time of current time step, used by TLM and cost integral */
     23:   PetscReal    time_step0;       /* Backup for ts->timestep, the step size of current time step, used by TLM and cost integral*/

     25:   /* context for sensitivity analysis */
     26:   PetscInt num_tlm;               /* Total number of tangent linear equations */
     27:   Vec     *VecsDeltaLam;          /* Increment of the adjoint sensitivity w.r.t IC at stage */
     28:   Vec     *VecsDeltaMu;           /* Increment of the adjoint sensitivity w.r.t P at stage */
     29:   Vec     *VecsSensiTemp;         /* Vector to be multiplied with Jacobian transpose */
     30:   Mat      MatFwdStages[2];       /* TLM Stages */
     31:   Mat      MatDeltaFwdSensip;     /* Increment of the forward sensitivity at stage */
     32:   Vec      VecDeltaFwdSensipCol;  /* Working vector for holding one column of the sensitivity matrix */
     33:   Mat      MatFwdSensip0;         /* backup for roll-backs due to events */
     34:   Mat      MatIntegralSensipTemp; /* Working vector for forward integral sensitivity */
     35:   Mat      MatIntegralSensip0;    /* backup for roll-backs due to events */
     36:   Vec     *VecsDeltaLam2;         /* Increment of the 2nd-order adjoint sensitivity w.r.t IC at stage */
     37:   Vec     *VecsDeltaMu2;          /* Increment of the 2nd-order adjoint sensitivity w.r.t P at stage */
     38:   Vec     *VecsSensi2Temp;        /* Working vectors that holds the residual for the second-order adjoint */
     39:   Vec     *VecsAffine;            /* Working vectors to store residuals */
     40:   /* context for error estimation */
     41:   Vec vec_sol_prev;
     42:   Vec vec_lte_work;
     43: } TS_Theta;

     45: static PetscErrorCode TSThetaGetX0AndXdot(TS ts, DM dm, Vec *X0, Vec *Xdot)
     46: {
     47:   TS_Theta *th = (TS_Theta *)ts->data;

     49:   PetscFunctionBegin;
     50:   if (X0) {
     51:     if (dm && dm != ts->dm) PetscCall(DMGetNamedGlobalVector(dm, "TSTheta_X0", X0));
     52:     else *X0 = ts->vec_sol;
     53:   }
     54:   if (Xdot) {
     55:     if (dm && dm != ts->dm) PetscCall(DMGetNamedGlobalVector(dm, "TSTheta_Xdot", Xdot));
     56:     else *Xdot = th->Xdot;
     57:   }
     58:   PetscFunctionReturn(PETSC_SUCCESS);
     59: }

     61: static PetscErrorCode TSThetaRestoreX0AndXdot(TS ts, DM dm, Vec *X0, Vec *Xdot)
     62: {
     63:   PetscFunctionBegin;
     64:   if (X0) {
     65:     if (dm && dm != ts->dm) PetscCall(DMRestoreNamedGlobalVector(dm, "TSTheta_X0", X0));
     66:   }
     67:   if (Xdot) {
     68:     if (dm && dm != ts->dm) PetscCall(DMRestoreNamedGlobalVector(dm, "TSTheta_Xdot", Xdot));
     69:   }
     70:   PetscFunctionReturn(PETSC_SUCCESS);
     71: }

     73: static PetscErrorCode DMCoarsenHook_TSTheta(DM fine, DM coarse, void *ctx)
     74: {
     75:   PetscFunctionBegin;
     76:   PetscFunctionReturn(PETSC_SUCCESS);
     77: }

     79: static PetscErrorCode DMRestrictHook_TSTheta(DM fine, Mat restrct, Vec rscale, Mat inject, DM coarse, void *ctx)
     80: {
     81:   TS  ts = (TS)ctx;
     82:   Vec X0, Xdot, X0_c, Xdot_c;

     84:   PetscFunctionBegin;
     85:   PetscCall(TSThetaGetX0AndXdot(ts, fine, &X0, &Xdot));
     86:   PetscCall(TSThetaGetX0AndXdot(ts, coarse, &X0_c, &Xdot_c));
     87:   PetscCall(MatRestrict(restrct, X0, X0_c));
     88:   PetscCall(MatRestrict(restrct, Xdot, Xdot_c));
     89:   PetscCall(VecPointwiseMult(X0_c, rscale, X0_c));
     90:   PetscCall(VecPointwiseMult(Xdot_c, rscale, Xdot_c));
     91:   PetscCall(TSThetaRestoreX0AndXdot(ts, fine, &X0, &Xdot));
     92:   PetscCall(TSThetaRestoreX0AndXdot(ts, coarse, &X0_c, &Xdot_c));
     93:   PetscFunctionReturn(PETSC_SUCCESS);
     94: }

     96: static PetscErrorCode DMSubDomainHook_TSTheta(DM dm, DM subdm, void *ctx)
     97: {
     98:   PetscFunctionBegin;
     99:   PetscFunctionReturn(PETSC_SUCCESS);
    100: }

    102: static PetscErrorCode DMSubDomainRestrictHook_TSTheta(DM dm, VecScatter gscat, VecScatter lscat, DM subdm, void *ctx)
    103: {
    104:   TS  ts = (TS)ctx;
    105:   Vec X0, Xdot, X0_sub, Xdot_sub;

    107:   PetscFunctionBegin;
    108:   PetscCall(TSThetaGetX0AndXdot(ts, dm, &X0, &Xdot));
    109:   PetscCall(TSThetaGetX0AndXdot(ts, subdm, &X0_sub, &Xdot_sub));

    111:   PetscCall(VecScatterBegin(gscat, X0, X0_sub, INSERT_VALUES, SCATTER_FORWARD));
    112:   PetscCall(VecScatterEnd(gscat, X0, X0_sub, INSERT_VALUES, SCATTER_FORWARD));

    114:   PetscCall(VecScatterBegin(gscat, Xdot, Xdot_sub, INSERT_VALUES, SCATTER_FORWARD));
    115:   PetscCall(VecScatterEnd(gscat, Xdot, Xdot_sub, INSERT_VALUES, SCATTER_FORWARD));

    117:   PetscCall(TSThetaRestoreX0AndXdot(ts, dm, &X0, &Xdot));
    118:   PetscCall(TSThetaRestoreX0AndXdot(ts, subdm, &X0_sub, &Xdot_sub));
    119:   PetscFunctionReturn(PETSC_SUCCESS);
    120: }

    122: static PetscErrorCode TSThetaEvaluateCostIntegral(TS ts)
    123: {
    124:   TS_Theta *th     = (TS_Theta *)ts->data;
    125:   TS        quadts = ts->quadraturets;

    127:   PetscFunctionBegin;
    128:   if (th->endpoint) {
    129:     /* Evolve ts->vec_costintegral to compute integrals */
    130:     if (th->Theta != 1.0) {
    131:       PetscCall(TSComputeRHSFunction(quadts, th->ptime0, th->X0, ts->vec_costintegrand));
    132:       PetscCall(VecAXPY(quadts->vec_sol, th->time_step0 * (1.0 - th->Theta), ts->vec_costintegrand));
    133:     }
    134:     PetscCall(TSComputeRHSFunction(quadts, ts->ptime, ts->vec_sol, ts->vec_costintegrand));
    135:     PetscCall(VecAXPY(quadts->vec_sol, th->time_step0 * th->Theta, ts->vec_costintegrand));
    136:   } else {
    137:     PetscCall(TSComputeRHSFunction(quadts, th->stage_time, th->X, ts->vec_costintegrand));
    138:     PetscCall(VecAXPY(quadts->vec_sol, th->time_step0, ts->vec_costintegrand));
    139:   }
    140:   PetscFunctionReturn(PETSC_SUCCESS);
    141: }

    143: static PetscErrorCode TSForwardCostIntegral_Theta(TS ts)
    144: {
    145:   TS_Theta *th     = (TS_Theta *)ts->data;
    146:   TS        quadts = ts->quadraturets;

    148:   PetscFunctionBegin;
    149:   /* backup cost integral */
    150:   PetscCall(VecCopy(quadts->vec_sol, th->VecCostIntegral0));
    151:   PetscCall(TSThetaEvaluateCostIntegral(ts));
    152:   PetscFunctionReturn(PETSC_SUCCESS);
    153: }

    155: static PetscErrorCode TSAdjointCostIntegral_Theta(TS ts)
    156: {
    157:   TS_Theta *th = (TS_Theta *)ts->data;

    159:   PetscFunctionBegin;
    160:   /* Like TSForwardCostIntegral(), the adjoint cost integral evaluation relies on ptime0 and time_step0. */
    161:   th->ptime0     = ts->ptime + ts->time_step;
    162:   th->time_step0 = -ts->time_step;
    163:   PetscCall(TSThetaEvaluateCostIntegral(ts));
    164:   PetscFunctionReturn(PETSC_SUCCESS);
    165: }

    167: static PetscErrorCode TSTheta_SNESSolve(TS ts, Vec b, Vec x)
    168: {
    169:   PetscInt nits, lits;

    171:   PetscFunctionBegin;
    172:   PetscCall(SNESSolve(ts->snes, b, x));
    173:   PetscCall(SNESGetIterationNumber(ts->snes, &nits));
    174:   PetscCall(SNESGetLinearSolveIterations(ts->snes, &lits));
    175:   ts->snes_its += nits;
    176:   ts->ksp_its += lits;
    177:   PetscFunctionReturn(PETSC_SUCCESS);
    178: }

    180: static PetscErrorCode TSResizeRegister_Theta(TS ts, PetscBool reg)
    181: {
    182:   TS_Theta *th = (TS_Theta *)ts->data;

    184:   PetscFunctionBegin;
    185:   if (reg) {
    186:     PetscCall(TSResizeRegisterVec(ts, "ts:theta:sol_prev", th->vec_sol_prev));
    187:     PetscCall(TSResizeRegisterVec(ts, "ts:theta:X0", th->X0));
    188:   } else {
    189:     PetscCall(TSResizeRetrieveVec(ts, "ts:theta:sol_prev", &th->vec_sol_prev));
    190:     PetscCall(PetscObjectReference((PetscObject)th->vec_sol_prev));
    191:     PetscCall(TSResizeRetrieveVec(ts, "ts:theta:X0", &th->X0));
    192:     PetscCall(PetscObjectReference((PetscObject)th->X0));
    193:   }
    194:   PetscFunctionReturn(PETSC_SUCCESS);
    195: }

    197: static PetscErrorCode TSStep_Theta(TS ts)
    198: {
    199:   TS_Theta *th         = (TS_Theta *)ts->data;
    200:   PetscInt  rejections = 0;
    201:   PetscBool stageok, accept = PETSC_TRUE;
    202:   PetscReal next_time_step = ts->time_step;

    204:   PetscFunctionBegin;
    205:   if (!ts->steprollback) {
    206:     if (th->vec_sol_prev) PetscCall(VecCopy(th->X0, th->vec_sol_prev));
    207:     PetscCall(VecCopy(ts->vec_sol, th->X0));
    208:   }

    210:   th->status = TS_STEP_INCOMPLETE;
    211:   while (!ts->reason && th->status != TS_STEP_COMPLETE) {
    212:     th->shift      = 1 / (th->Theta * ts->time_step);
    213:     th->stage_time = ts->ptime + (th->endpoint ? (PetscReal)1 : th->Theta) * ts->time_step;
    214:     PetscCall(VecCopy(th->X0, th->X));
    215:     if (th->extrapolate && !ts->steprestart) PetscCall(VecAXPY(th->X, 1 / th->shift, th->Xdot));
    216:     if (th->endpoint) { /* This formulation assumes linear time-independent mass matrix */
    217:       if (!th->affine) PetscCall(VecDuplicate(ts->vec_sol, &th->affine));
    218:       PetscCall(VecZeroEntries(th->Xdot));
    219:       PetscCall(TSComputeIFunction(ts, ts->ptime, th->X0, th->Xdot, th->affine, PETSC_FALSE));
    220:       PetscCall(VecScale(th->affine, (th->Theta - 1) / th->Theta));
    221:     }
    222:     PetscCall(TSPreStage(ts, th->stage_time));
    223:     PetscCall(TSTheta_SNESSolve(ts, th->endpoint ? th->affine : NULL, th->X));
    224:     PetscCall(TSPostStage(ts, th->stage_time, 0, &th->X));
    225:     PetscCall(TSAdaptCheckStage(ts->adapt, ts, th->stage_time, th->X, &stageok));
    226:     if (!stageok) goto reject_step;

    228:     if (th->endpoint) {
    229:       PetscCall(VecCopy(th->X, ts->vec_sol));
    230:     } else {
    231:       PetscCall(VecAXPBYPCZ(th->Xdot, -th->shift, th->shift, 0, th->X0, th->X)); /* th->Xdot is needed by TSInterpolate_Theta */
    232:       if (th->Theta == 1.0) PetscCall(VecCopy(th->X, ts->vec_sol));              /* BEULER, stage already checked */
    233:       else {
    234:         PetscCall(VecAXPY(ts->vec_sol, ts->time_step, th->Xdot));
    235:         PetscCall(TSAdaptCheckStage(ts->adapt, ts, ts->ptime + ts->time_step, ts->vec_sol, &stageok));
    236:         if (!stageok) {
    237:           PetscCall(VecCopy(th->X0, ts->vec_sol));
    238:           goto reject_step;
    239:         }
    240:       }
    241:     }

    243:     th->status = TS_STEP_PENDING;
    244:     PetscCall(TSAdaptChoose(ts->adapt, ts, ts->time_step, NULL, &next_time_step, &accept));
    245:     th->status = accept ? TS_STEP_COMPLETE : TS_STEP_INCOMPLETE;
    246:     if (!accept) {
    247:       PetscCall(VecCopy(th->X0, ts->vec_sol));
    248:       ts->time_step = next_time_step;
    249:       goto reject_step;
    250:     }

    252:     if (ts->forward_solve || ts->costintegralfwd) { /* Save the info for the later use in cost integral evaluation */
    253:       th->ptime0     = ts->ptime;
    254:       th->time_step0 = ts->time_step;
    255:     }
    256:     ts->ptime += ts->time_step;
    257:     ts->time_step = next_time_step;
    258:     break;

    260:   reject_step:
    261:     ts->reject++;
    262:     accept = PETSC_FALSE;
    263:     if (!ts->reason && ++rejections > ts->max_reject && ts->max_reject >= 0) {
    264:       ts->reason = TS_DIVERGED_STEP_REJECTED;
    265:       PetscCall(PetscInfo(ts, "Step=%" PetscInt_FMT ", step rejections %" PetscInt_FMT " greater than current TS allowed, stopping solve\n", ts->steps, rejections));
    266:     }
    267:   }
    268:   PetscFunctionReturn(PETSC_SUCCESS);
    269: }

    271: static PetscErrorCode TSAdjointStepBEuler_Private(TS ts)
    272: {
    273:   TS_Theta      *th           = (TS_Theta *)ts->data;
    274:   TS             quadts       = ts->quadraturets;
    275:   Vec           *VecsDeltaLam = th->VecsDeltaLam, *VecsDeltaMu = th->VecsDeltaMu, *VecsSensiTemp = th->VecsSensiTemp;
    276:   Vec           *VecsDeltaLam2 = th->VecsDeltaLam2, *VecsDeltaMu2 = th->VecsDeltaMu2, *VecsSensi2Temp = th->VecsSensi2Temp;
    277:   PetscInt       nadj;
    278:   Mat            J, Jpre, quadJ = NULL, quadJp = NULL;
    279:   KSP            ksp;
    280:   PetscScalar   *xarr;
    281:   TSEquationType eqtype;
    282:   PetscBool      isexplicitode = PETSC_FALSE;
    283:   PetscReal      adjoint_time_step;

    285:   PetscFunctionBegin;
    286:   PetscCall(TSGetEquationType(ts, &eqtype));
    287:   if (eqtype == TS_EQ_ODE_EXPLICIT) {
    288:     isexplicitode = PETSC_TRUE;
    289:     VecsDeltaLam  = ts->vecs_sensi;
    290:     VecsDeltaLam2 = ts->vecs_sensi2;
    291:   }
    292:   th->status = TS_STEP_INCOMPLETE;
    293:   PetscCall(SNESGetKSP(ts->snes, &ksp));
    294:   PetscCall(TSGetIJacobian(ts, &J, &Jpre, NULL, NULL));
    295:   if (quadts) {
    296:     PetscCall(TSGetRHSJacobian(quadts, &quadJ, NULL, NULL, NULL));
    297:     PetscCall(TSGetRHSJacobianP(quadts, &quadJp, NULL, NULL));
    298:   }

    300:   th->stage_time    = ts->ptime;
    301:   adjoint_time_step = -ts->time_step; /* always positive since time_step is negative */

    303:   /* Build RHS for first-order adjoint lambda_{n+1}/h + r_u^T(n+1) */
    304:   if (quadts) PetscCall(TSComputeRHSJacobian(quadts, th->stage_time, ts->vec_sol, quadJ, NULL));

    306:   for (nadj = 0; nadj < ts->numcost; nadj++) {
    307:     PetscCall(VecCopy(ts->vecs_sensi[nadj], VecsSensiTemp[nadj]));
    308:     PetscCall(VecScale(VecsSensiTemp[nadj], 1. / adjoint_time_step)); /* lambda_{n+1}/h */
    309:     if (quadJ) {
    310:       PetscCall(MatDenseGetColumn(quadJ, nadj, &xarr));
    311:       PetscCall(VecPlaceArray(ts->vec_drdu_col, xarr));
    312:       PetscCall(VecAXPY(VecsSensiTemp[nadj], 1., ts->vec_drdu_col));
    313:       PetscCall(VecResetArray(ts->vec_drdu_col));
    314:       PetscCall(MatDenseRestoreColumn(quadJ, &xarr));
    315:     }
    316:   }

    318:   /* Build LHS for first-order adjoint */
    319:   th->shift = 1. / adjoint_time_step;
    320:   PetscCall(TSComputeSNESJacobian(ts, ts->vec_sol, J, Jpre));
    321:   PetscCall(KSPSetOperators(ksp, J, Jpre));

    323:   /* Solve stage equation LHS*lambda_s = RHS for first-order adjoint */
    324:   for (nadj = 0; nadj < ts->numcost; nadj++) {
    325:     KSPConvergedReason kspreason;
    326:     PetscCall(KSPSolveTranspose(ksp, VecsSensiTemp[nadj], VecsDeltaLam[nadj]));
    327:     PetscCall(KSPGetConvergedReason(ksp, &kspreason));
    328:     if (kspreason < 0) {
    329:       ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
    330:       PetscCall(PetscInfo(ts, "Step=%" PetscInt_FMT ", %" PetscInt_FMT "th cost function, transposed linear solve fails, stopping 1st-order adjoint solve\n", ts->steps, nadj));
    331:     }
    332:   }

    334:   if (ts->vecs_sensi2) { /* U_{n+1} */
    335:     /* Get w1 at t_{n+1} from TLM matrix */
    336:     PetscCall(MatDenseGetColumn(ts->mat_sensip, 0, &xarr));
    337:     PetscCall(VecPlaceArray(ts->vec_sensip_col, xarr));
    338:     /* lambda_s^T F_UU w_1 */
    339:     PetscCall(TSComputeIHessianProductFunctionUU(ts, th->stage_time, ts->vec_sol, VecsDeltaLam, ts->vec_sensip_col, ts->vecs_fuu));
    340:     /* lambda_s^T F_UP w_2 */
    341:     PetscCall(TSComputeIHessianProductFunctionUP(ts, th->stage_time, ts->vec_sol, VecsDeltaLam, ts->vec_dir, ts->vecs_fup));
    342:     for (nadj = 0; nadj < ts->numcost; nadj++) { /* compute the residual */
    343:       PetscCall(VecCopy(ts->vecs_sensi2[nadj], VecsSensi2Temp[nadj]));
    344:       PetscCall(VecScale(VecsSensi2Temp[nadj], 1. / adjoint_time_step));
    345:       PetscCall(VecAXPY(VecsSensi2Temp[nadj], -1., ts->vecs_fuu[nadj]));
    346:       if (ts->vecs_fup) PetscCall(VecAXPY(VecsSensi2Temp[nadj], -1., ts->vecs_fup[nadj]));
    347:     }
    348:     /* Solve stage equation LHS X = RHS for second-order adjoint */
    349:     for (nadj = 0; nadj < ts->numcost; nadj++) {
    350:       KSPConvergedReason kspreason;
    351:       PetscCall(KSPSolveTranspose(ksp, VecsSensi2Temp[nadj], VecsDeltaLam2[nadj]));
    352:       PetscCall(KSPGetConvergedReason(ksp, &kspreason));
    353:       if (kspreason < 0) {
    354:         ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
    355:         PetscCall(PetscInfo(ts, "Step=%" PetscInt_FMT ", %" PetscInt_FMT "th cost function, transposed linear solve fails, stopping 2nd-order adjoint solve\n", ts->steps, nadj));
    356:       }
    357:     }
    358:   }

    360:   /* Update sensitivities, and evaluate integrals if there is any */
    361:   if (!isexplicitode) {
    362:     th->shift = 0.0;
    363:     PetscCall(TSComputeSNESJacobian(ts, ts->vec_sol, J, Jpre));
    364:     PetscCall(KSPSetOperators(ksp, J, Jpre));
    365:     for (nadj = 0; nadj < ts->numcost; nadj++) {
    366:       /* Add f_U \lambda_s to the original RHS */
    367:       PetscCall(VecScale(VecsSensiTemp[nadj], -1.));
    368:       PetscCall(MatMultTransposeAdd(J, VecsDeltaLam[nadj], VecsSensiTemp[nadj], VecsSensiTemp[nadj]));
    369:       PetscCall(VecScale(VecsSensiTemp[nadj], -adjoint_time_step));
    370:       PetscCall(VecCopy(VecsSensiTemp[nadj], ts->vecs_sensi[nadj]));
    371:       if (ts->vecs_sensi2) {
    372:         PetscCall(MatMultTransposeAdd(J, VecsDeltaLam2[nadj], VecsSensi2Temp[nadj], VecsSensi2Temp[nadj]));
    373:         PetscCall(VecScale(VecsSensi2Temp[nadj], -adjoint_time_step));
    374:         PetscCall(VecCopy(VecsSensi2Temp[nadj], ts->vecs_sensi2[nadj]));
    375:       }
    376:     }
    377:   }
    378:   if (ts->vecs_sensip) {
    379:     PetscCall(VecAXPBYPCZ(th->Xdot, -1. / adjoint_time_step, 1.0 / adjoint_time_step, 0, th->X0, ts->vec_sol));
    380:     PetscCall(TSComputeIJacobianP(ts, th->stage_time, ts->vec_sol, th->Xdot, 1. / adjoint_time_step, ts->Jacp, PETSC_FALSE)); /* get -f_p */
    381:     if (quadts) PetscCall(TSComputeRHSJacobianP(quadts, th->stage_time, ts->vec_sol, quadJp));
    382:     if (ts->vecs_sensi2p) {
    383:       /* lambda_s^T F_PU w_1 */
    384:       PetscCall(TSComputeIHessianProductFunctionPU(ts, th->stage_time, ts->vec_sol, VecsDeltaLam, ts->vec_sensip_col, ts->vecs_fpu));
    385:       /* lambda_s^T F_PP w_2 */
    386:       PetscCall(TSComputeIHessianProductFunctionPP(ts, th->stage_time, ts->vec_sol, VecsDeltaLam, ts->vec_dir, ts->vecs_fpp));
    387:     }

    389:     for (nadj = 0; nadj < ts->numcost; nadj++) {
    390:       PetscCall(MatMultTranspose(ts->Jacp, VecsDeltaLam[nadj], VecsDeltaMu[nadj]));
    391:       PetscCall(VecAXPY(ts->vecs_sensip[nadj], -adjoint_time_step, VecsDeltaMu[nadj]));
    392:       if (quadJp) {
    393:         PetscCall(MatDenseGetColumn(quadJp, nadj, &xarr));
    394:         PetscCall(VecPlaceArray(ts->vec_drdp_col, xarr));
    395:         PetscCall(VecAXPY(ts->vecs_sensip[nadj], adjoint_time_step, ts->vec_drdp_col));
    396:         PetscCall(VecResetArray(ts->vec_drdp_col));
    397:         PetscCall(MatDenseRestoreColumn(quadJp, &xarr));
    398:       }
    399:       if (ts->vecs_sensi2p) {
    400:         PetscCall(MatMultTranspose(ts->Jacp, VecsDeltaLam2[nadj], VecsDeltaMu2[nadj]));
    401:         PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step, VecsDeltaMu2[nadj]));
    402:         if (ts->vecs_fpu) PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step, ts->vecs_fpu[nadj]));
    403:         if (ts->vecs_fpp) PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step, ts->vecs_fpp[nadj]));
    404:       }
    405:     }
    406:   }

    408:   if (ts->vecs_sensi2) {
    409:     PetscCall(VecResetArray(ts->vec_sensip_col));
    410:     PetscCall(MatDenseRestoreColumn(ts->mat_sensip, &xarr));
    411:   }
    412:   th->status = TS_STEP_COMPLETE;
    413:   PetscFunctionReturn(PETSC_SUCCESS);
    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, usually ts->ptime is larger than adjoint_ptime) */

    429:   PetscFunctionBegin;
    430:   if (th->Theta == 1.) {
    431:     PetscCall(TSAdjointStepBEuler_Private(ts));
    432:     PetscFunctionReturn(PETSC_SUCCESS);
    433:   }
    434:   th->status = TS_STEP_INCOMPLETE;
    435:   PetscCall(SNESGetKSP(ts->snes, &ksp));
    436:   PetscCall(TSGetIJacobian(ts, &J, &Jpre, NULL, NULL));
    437:   if (quadts) {
    438:     PetscCall(TSGetRHSJacobian(quadts, &quadJ, NULL, NULL, NULL));
    439:     PetscCall(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:   if (!th->endpoint) {
    447:     /* recover th->X0 using vec_sol and the stage value th->X */
    448:     PetscCall(VecAXPBYPCZ(th->X0, 1.0 / (1.0 - th->Theta), th->Theta / (th->Theta - 1.0), 0, th->X, ts->vec_sol));
    449:   }

    451:   /* Build RHS for first-order adjoint */
    452:   /* Cost function has an integral term */
    453:   if (quadts) {
    454:     if (th->endpoint) {
    455:       PetscCall(TSComputeRHSJacobian(quadts, th->stage_time, ts->vec_sol, quadJ, NULL));
    456:     } else {
    457:       PetscCall(TSComputeRHSJacobian(quadts, th->stage_time, th->X, quadJ, NULL));
    458:     }
    459:   }

    461:   for (nadj = 0; nadj < ts->numcost; nadj++) {
    462:     PetscCall(VecCopy(ts->vecs_sensi[nadj], VecsSensiTemp[nadj]));
    463:     PetscCall(VecScale(VecsSensiTemp[nadj], 1. / (th->Theta * adjoint_time_step)));
    464:     if (quadJ) {
    465:       PetscCall(MatDenseGetColumn(quadJ, nadj, &xarr));
    466:       PetscCall(VecPlaceArray(ts->vec_drdu_col, xarr));
    467:       PetscCall(VecAXPY(VecsSensiTemp[nadj], 1., ts->vec_drdu_col));
    468:       PetscCall(VecResetArray(ts->vec_drdu_col));
    469:       PetscCall(MatDenseRestoreColumn(quadJ, &xarr));
    470:     }
    471:   }

    473:   /* Build LHS for first-order adjoint */
    474:   th->shift = 1. / (th->Theta * adjoint_time_step);
    475:   if (th->endpoint) {
    476:     PetscCall(TSComputeSNESJacobian(ts, ts->vec_sol, J, Jpre));
    477:   } else {
    478:     PetscCall(TSComputeSNESJacobian(ts, th->X, J, Jpre));
    479:   }
    480:   PetscCall(KSPSetOperators(ksp, J, Jpre));

    482:   /* Solve stage equation LHS*lambda_s = RHS for first-order adjoint */
    483:   for (nadj = 0; nadj < ts->numcost; nadj++) {
    484:     KSPConvergedReason kspreason;
    485:     PetscCall(KSPSolveTranspose(ksp, VecsSensiTemp[nadj], VecsDeltaLam[nadj]));
    486:     PetscCall(KSPGetConvergedReason(ksp, &kspreason));
    487:     if (kspreason < 0) {
    488:       ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
    489:       PetscCall(PetscInfo(ts, "Step=%" PetscInt_FMT ", %" PetscInt_FMT "th cost function, transposed linear solve fails, stopping 1st-order adjoint solve\n", ts->steps, nadj));
    490:     }
    491:   }

    493:   /* Second-order adjoint */
    494:   if (ts->vecs_sensi2) { /* U_{n+1} */
    495:     PetscCheck(th->endpoint, PetscObjectComm((PetscObject)ts), PETSC_ERR_SUP, "Operation not implemented in TS_Theta");
    496:     /* Get w1 at t_{n+1} from TLM matrix */
    497:     PetscCall(MatDenseGetColumn(ts->mat_sensip, 0, &xarr));
    498:     PetscCall(VecPlaceArray(ts->vec_sensip_col, xarr));
    499:     /* lambda_s^T F_UU w_1 */
    500:     PetscCall(TSComputeIHessianProductFunctionUU(ts, th->stage_time, ts->vec_sol, VecsDeltaLam, ts->vec_sensip_col, ts->vecs_fuu));
    501:     PetscCall(VecResetArray(ts->vec_sensip_col));
    502:     PetscCall(MatDenseRestoreColumn(ts->mat_sensip, &xarr));
    503:     /* lambda_s^T F_UP w_2 */
    504:     PetscCall(TSComputeIHessianProductFunctionUP(ts, th->stage_time, ts->vec_sol, VecsDeltaLam, ts->vec_dir, ts->vecs_fup));
    505:     for (nadj = 0; nadj < ts->numcost; nadj++) { /* compute the residual */
    506:       PetscCall(VecCopy(ts->vecs_sensi2[nadj], VecsSensi2Temp[nadj]));
    507:       PetscCall(VecScale(VecsSensi2Temp[nadj], th->shift));
    508:       PetscCall(VecAXPY(VecsSensi2Temp[nadj], -1., ts->vecs_fuu[nadj]));
    509:       if (ts->vecs_fup) PetscCall(VecAXPY(VecsSensi2Temp[nadj], -1., ts->vecs_fup[nadj]));
    510:     }
    511:     /* Solve stage equation LHS X = RHS for second-order adjoint */
    512:     for (nadj = 0; nadj < ts->numcost; nadj++) {
    513:       KSPConvergedReason kspreason;
    514:       PetscCall(KSPSolveTranspose(ksp, VecsSensi2Temp[nadj], VecsDeltaLam2[nadj]));
    515:       PetscCall(KSPGetConvergedReason(ksp, &kspreason));
    516:       if (kspreason < 0) {
    517:         ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
    518:         PetscCall(PetscInfo(ts, "Step=%" PetscInt_FMT ", %" PetscInt_FMT "th cost function, transposed linear solve fails, stopping 2nd-order adjoint solve\n", ts->steps, nadj));
    519:       }
    520:     }
    521:   }

    523:   /* Update sensitivities, and evaluate integrals if there is any */
    524:   if (th->endpoint) { /* two-stage Theta methods with th->Theta!=1, th->Theta==1 leads to BEuler */
    525:     th->shift      = 1. / ((th->Theta - 1.) * adjoint_time_step);
    526:     th->stage_time = adjoint_ptime;
    527:     PetscCall(TSComputeSNESJacobian(ts, th->X0, J, Jpre));
    528:     PetscCall(KSPSetOperators(ksp, J, Jpre));
    529:     /* R_U at t_n */
    530:     if (quadts) PetscCall(TSComputeRHSJacobian(quadts, adjoint_ptime, th->X0, quadJ, NULL));
    531:     for (nadj = 0; nadj < ts->numcost; nadj++) {
    532:       PetscCall(MatMultTranspose(J, VecsDeltaLam[nadj], ts->vecs_sensi[nadj]));
    533:       if (quadJ) {
    534:         PetscCall(MatDenseGetColumn(quadJ, nadj, &xarr));
    535:         PetscCall(VecPlaceArray(ts->vec_drdu_col, xarr));
    536:         PetscCall(VecAXPY(ts->vecs_sensi[nadj], -1., ts->vec_drdu_col));
    537:         PetscCall(VecResetArray(ts->vec_drdu_col));
    538:         PetscCall(MatDenseRestoreColumn(quadJ, &xarr));
    539:       }
    540:       PetscCall(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:       PetscCall(MatDenseGetColumn(th->MatFwdSensip0, 0, &xarr));
    547:       PetscCall(VecPlaceArray(ts->vec_sensip_col, xarr));
    548:       /* lambda_s^T F_UU w_1 */
    549:       PetscCall(TSComputeIHessianProductFunctionUU(ts, adjoint_ptime, th->X0, VecsDeltaLam, ts->vec_sensip_col, ts->vecs_fuu));
    550:       PetscCall(VecResetArray(ts->vec_sensip_col));
    551:       PetscCall(MatDenseRestoreColumn(th->MatFwdSensip0, &xarr));
    552:       /* lambda_s^T F_UU w_2 */
    553:       PetscCall(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:         PetscCall(MatMultTranspose(J, VecsDeltaLam2[nadj], ts->vecs_sensi2[nadj]));
    557:         PetscCall(VecAXPY(ts->vecs_sensi2[nadj], 1., ts->vecs_fuu[nadj]));
    558:         if (ts->vecs_fup) PetscCall(VecAXPY(ts->vecs_sensi2[nadj], 1., ts->vecs_fup[nadj]));
    559:         PetscCall(VecScale(ts->vecs_sensi2[nadj], 1. / th->shift));
    560:       }
    561:     }

    563:     th->stage_time = ts->ptime; /* recover the old value */

    565:     if (ts->vecs_sensip) { /* sensitivities wrt parameters */
    566:       /* U_{n+1} */
    567:       th->shift = 1.0 / (adjoint_time_step * th->Theta);
    568:       PetscCall(VecAXPBYPCZ(th->Xdot, -th->shift, th->shift, 0, th->X0, ts->vec_sol));
    569:       PetscCall(TSComputeIJacobianP(ts, th->stage_time, ts->vec_sol, th->Xdot, -1. / (th->Theta * adjoint_time_step), ts->Jacp, PETSC_FALSE));
    570:       if (quadts) PetscCall(TSComputeRHSJacobianP(quadts, th->stage_time, ts->vec_sol, quadJp));
    571:       for (nadj = 0; nadj < ts->numcost; nadj++) {
    572:         PetscCall(MatMultTranspose(ts->Jacp, VecsDeltaLam[nadj], VecsDeltaMu[nadj]));
    573:         PetscCall(VecAXPY(ts->vecs_sensip[nadj], -adjoint_time_step * th->Theta, VecsDeltaMu[nadj]));
    574:         if (quadJp) {
    575:           PetscCall(MatDenseGetColumn(quadJp, nadj, &xarr));
    576:           PetscCall(VecPlaceArray(ts->vec_drdp_col, xarr));
    577:           PetscCall(VecAXPY(ts->vecs_sensip[nadj], adjoint_time_step * th->Theta, ts->vec_drdp_col));
    578:           PetscCall(VecResetArray(ts->vec_drdp_col));
    579:           PetscCall(MatDenseRestoreColumn(quadJp, &xarr));
    580:         }
    581:       }
    582:       if (ts->vecs_sensi2p) { /* second-order */
    583:         /* Get w1 at t_{n+1} from TLM matrix */
    584:         PetscCall(MatDenseGetColumn(ts->mat_sensip, 0, &xarr));
    585:         PetscCall(VecPlaceArray(ts->vec_sensip_col, xarr));
    586:         /* lambda_s^T F_PU w_1 */
    587:         PetscCall(TSComputeIHessianProductFunctionPU(ts, th->stage_time, ts->vec_sol, VecsDeltaLam, ts->vec_sensip_col, ts->vecs_fpu));
    588:         PetscCall(VecResetArray(ts->vec_sensip_col));
    589:         PetscCall(MatDenseRestoreColumn(ts->mat_sensip, &xarr));

    591:         /* lambda_s^T F_PP w_2 */
    592:         PetscCall(TSComputeIHessianProductFunctionPP(ts, th->stage_time, ts->vec_sol, VecsDeltaLam, ts->vec_dir, ts->vecs_fpp));
    593:         for (nadj = 0; nadj < ts->numcost; nadj++) {
    594:           /* 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)  */
    595:           PetscCall(MatMultTranspose(ts->Jacp, VecsDeltaLam2[nadj], VecsDeltaMu2[nadj]));
    596:           PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step * th->Theta, VecsDeltaMu2[nadj]));
    597:           if (ts->vecs_fpu) PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step * th->Theta, ts->vecs_fpu[nadj]));
    598:           if (ts->vecs_fpp) PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step * th->Theta, ts->vecs_fpp[nadj]));
    599:         }
    600:       }

    602:       /* U_s */
    603:       PetscCall(VecZeroEntries(th->Xdot));
    604:       PetscCall(TSComputeIJacobianP(ts, adjoint_ptime, th->X0, th->Xdot, 1. / ((th->Theta - 1.0) * adjoint_time_step), ts->Jacp, PETSC_FALSE));
    605:       if (quadts) PetscCall(TSComputeRHSJacobianP(quadts, adjoint_ptime, th->X0, quadJp));
    606:       for (nadj = 0; nadj < ts->numcost; nadj++) {
    607:         PetscCall(MatMultTranspose(ts->Jacp, VecsDeltaLam[nadj], VecsDeltaMu[nadj]));
    608:         PetscCall(VecAXPY(ts->vecs_sensip[nadj], -adjoint_time_step * (1.0 - th->Theta), VecsDeltaMu[nadj]));
    609:         if (quadJp) {
    610:           PetscCall(MatDenseGetColumn(quadJp, nadj, &xarr));
    611:           PetscCall(VecPlaceArray(ts->vec_drdp_col, xarr));
    612:           PetscCall(VecAXPY(ts->vecs_sensip[nadj], adjoint_time_step * (1.0 - th->Theta), ts->vec_drdp_col));
    613:           PetscCall(VecResetArray(ts->vec_drdp_col));
    614:           PetscCall(MatDenseRestoreColumn(quadJp, &xarr));
    615:         }
    616:         if (ts->vecs_sensi2p) { /* second-order */
    617:           /* Get w1 at t_n from TLM matrix */
    618:           PetscCall(MatDenseGetColumn(th->MatFwdSensip0, 0, &xarr));
    619:           PetscCall(VecPlaceArray(ts->vec_sensip_col, xarr));
    620:           /* lambda_s^T F_PU w_1 */
    621:           PetscCall(TSComputeIHessianProductFunctionPU(ts, adjoint_ptime, th->X0, VecsDeltaLam, ts->vec_sensip_col, ts->vecs_fpu));
    622:           PetscCall(VecResetArray(ts->vec_sensip_col));
    623:           PetscCall(MatDenseRestoreColumn(th->MatFwdSensip0, &xarr));
    624:           /* lambda_s^T F_PP w_2 */
    625:           PetscCall(TSComputeIHessianProductFunctionPP(ts, adjoint_ptime, th->X0, VecsDeltaLam, ts->vec_dir, ts->vecs_fpp));
    626:           for (nadj = 0; nadj < ts->numcost; nadj++) {
    627:             /* 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) */
    628:             PetscCall(MatMultTranspose(ts->Jacp, VecsDeltaLam2[nadj], VecsDeltaMu2[nadj]));
    629:             PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step * (1.0 - th->Theta), VecsDeltaMu2[nadj]));
    630:             if (ts->vecs_fpu) PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step * (1.0 - th->Theta), ts->vecs_fpu[nadj]));
    631:             if (ts->vecs_fpp) PetscCall(VecAXPY(ts->vecs_sensi2p[nadj], -adjoint_time_step * (1.0 - th->Theta), ts->vecs_fpp[nadj]));
    632:           }
    633:         }
    634:       }
    635:     }
    636:   } else { /* one-stage case */
    637:     th->shift = 0.0;
    638:     PetscCall(TSComputeSNESJacobian(ts, th->X, J, Jpre)); /* get -f_y */
    639:     PetscCall(KSPSetOperators(ksp, J, Jpre));
    640:     if (quadts) PetscCall(TSComputeRHSJacobian(quadts, th->stage_time, th->X, quadJ, NULL));
    641:     for (nadj = 0; nadj < ts->numcost; nadj++) {
    642:       PetscCall(MatMultTranspose(J, VecsDeltaLam[nadj], VecsSensiTemp[nadj]));
    643:       PetscCall(VecAXPY(ts->vecs_sensi[nadj], -adjoint_time_step, VecsSensiTemp[nadj]));
    644:       if (quadJ) {
    645:         PetscCall(MatDenseGetColumn(quadJ, nadj, &xarr));
    646:         PetscCall(VecPlaceArray(ts->vec_drdu_col, xarr));
    647:         PetscCall(VecAXPY(ts->vecs_sensi[nadj], adjoint_time_step, ts->vec_drdu_col));
    648:         PetscCall(VecResetArray(ts->vec_drdu_col));
    649:         PetscCall(MatDenseRestoreColumn(quadJ, &xarr));
    650:       }
    651:     }
    652:     if (ts->vecs_sensip) {
    653:       th->shift = 1.0 / (adjoint_time_step * th->Theta);
    654:       PetscCall(VecAXPBYPCZ(th->Xdot, -th->shift, th->shift, 0, th->X0, th->X));
    655:       PetscCall(TSComputeIJacobianP(ts, th->stage_time, th->X, th->Xdot, th->shift, ts->Jacp, PETSC_FALSE));
    656:       if (quadts) PetscCall(TSComputeRHSJacobianP(quadts, th->stage_time, th->X, quadJp));
    657:       for (nadj = 0; nadj < ts->numcost; nadj++) {
    658:         PetscCall(MatMultTranspose(ts->Jacp, VecsDeltaLam[nadj], VecsDeltaMu[nadj]));
    659:         PetscCall(VecAXPY(ts->vecs_sensip[nadj], -adjoint_time_step, VecsDeltaMu[nadj]));
    660:         if (quadJp) {
    661:           PetscCall(MatDenseGetColumn(quadJp, nadj, &xarr));
    662:           PetscCall(VecPlaceArray(ts->vec_drdp_col, xarr));
    663:           PetscCall(VecAXPY(ts->vecs_sensip[nadj], adjoint_time_step, ts->vec_drdp_col));
    664:           PetscCall(VecResetArray(ts->vec_drdp_col));
    665:           PetscCall(MatDenseRestoreColumn(quadJp, &xarr));
    666:         }
    667:       }
    668:     }
    669:   }

    671:   th->status = TS_STEP_COMPLETE;
    672:   PetscFunctionReturn(PETSC_SUCCESS);
    673: }

    675: static PetscErrorCode TSInterpolate_Theta(TS ts, PetscReal t, Vec X)
    676: {
    677:   TS_Theta *th = (TS_Theta *)ts->data;
    678:   PetscReal dt = t - ts->ptime;

    680:   PetscFunctionBegin;
    681:   PetscCall(VecCopy(ts->vec_sol, th->X));
    682:   if (th->endpoint) dt *= th->Theta;
    683:   PetscCall(VecWAXPY(X, dt, th->Xdot, th->X));
    684:   PetscFunctionReturn(PETSC_SUCCESS);
    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;

    694:   PetscFunctionBegin;
    695:   if (!th->vec_sol_prev) {
    696:     *wlte = -1;
    697:     PetscFunctionReturn(PETSC_SUCCESS);
    698:   }
    699:   /* Cannot compute LTE in first step or in restart after event */
    700:   if (ts->steprestart) {
    701:     *wlte = -1;
    702:     PetscFunctionReturn(PETSC_SUCCESS);
    703:   }
    704:   /* Compute LTE using backward differences with non-constant time step */
    705:   {
    706:     PetscReal   h = ts->time_step, h_prev = ts->ptime - ts->ptime_prev;
    707:     PetscReal   a = 1 + h_prev / h;
    708:     PetscScalar scal[3];
    709:     Vec         vecs[3];

    711:     scal[0] = -1 / a;
    712:     scal[1] = +1 / (a - 1);
    713:     scal[2] = -1 / (a * (a - 1));
    714:     vecs[0] = X;
    715:     vecs[1] = th->X0;
    716:     vecs[2] = th->vec_sol_prev;
    717:     PetscCall(VecCopy(X, Y));
    718:     PetscCall(VecMAXPY(Y, 3, scal, vecs));
    719:     PetscCall(TSErrorWeightedNorm(ts, X, Y, wnormtype, wlte, &wltea, &wlter));
    720:   }
    721:   if (order) *order = 2;
    722:   PetscFunctionReturn(PETSC_SUCCESS);
    723: }

    725: static PetscErrorCode TSRollBack_Theta(TS ts)
    726: {
    727:   TS_Theta *th     = (TS_Theta *)ts->data;
    728:   TS        quadts = ts->quadraturets;

    730:   PetscFunctionBegin;
    731:   if (quadts && ts->costintegralfwd) PetscCall(VecCopy(th->VecCostIntegral0, quadts->vec_sol));
    732:   th->status = TS_STEP_INCOMPLETE;
    733:   if (ts->mat_sensip) PetscCall(MatCopy(th->MatFwdSensip0, ts->mat_sensip, SAME_NONZERO_PATTERN));
    734:   if (quadts && quadts->mat_sensip) PetscCall(MatCopy(th->MatIntegralSensip0, quadts->mat_sensip, SAME_NONZERO_PATTERN));
    735:   PetscFunctionReturn(PETSC_SUCCESS);
    736: }

    738: static PetscErrorCode TSForwardStep_Theta(TS ts)
    739: {
    740:   TS_Theta    *th                   = (TS_Theta *)ts->data;
    741:   TS           quadts               = ts->quadraturets;
    742:   Mat          MatDeltaFwdSensip    = th->MatDeltaFwdSensip;
    743:   Vec          VecDeltaFwdSensipCol = th->VecDeltaFwdSensipCol;
    744:   PetscInt     ntlm;
    745:   KSP          ksp;
    746:   Mat          J, Jpre, quadJ = NULL, quadJp = NULL;
    747:   PetscScalar *barr, *xarr;
    748:   PetscReal    previous_shift;

    750:   PetscFunctionBegin;
    751:   previous_shift = th->shift;
    752:   PetscCall(MatCopy(ts->mat_sensip, th->MatFwdSensip0, SAME_NONZERO_PATTERN));

    754:   if (quadts && quadts->mat_sensip) PetscCall(MatCopy(quadts->mat_sensip, th->MatIntegralSensip0, SAME_NONZERO_PATTERN));
    755:   PetscCall(SNESGetKSP(ts->snes, &ksp));
    756:   PetscCall(TSGetIJacobian(ts, &J, &Jpre, NULL, NULL));
    757:   if (quadts) {
    758:     PetscCall(TSGetRHSJacobian(quadts, &quadJ, NULL, NULL, NULL));
    759:     PetscCall(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:     PetscCall(TSComputeIJacobian(ts, th->ptime0, th->X0, th->Xdot, th->shift, J, Jpre, PETSC_FALSE));
    766:     PetscCall(MatMatMult(J, ts->mat_sensip, MAT_REUSE_MATRIX, PETSC_DETERMINE, &MatDeltaFwdSensip));
    767:     PetscCall(MatScale(MatDeltaFwdSensip, (th->Theta - 1.) / th->Theta));

    769:     /* Add the f_p forcing terms */
    770:     if (ts->Jacp) {
    771:       PetscCall(VecZeroEntries(th->Xdot));
    772:       PetscCall(TSComputeIJacobianP(ts, th->ptime0, th->X0, th->Xdot, th->shift, ts->Jacp, PETSC_FALSE));
    773:       PetscCall(MatAXPY(MatDeltaFwdSensip, (th->Theta - 1.) / th->Theta, ts->Jacp, SUBSET_NONZERO_PATTERN));
    774:       th->shift = previous_shift;
    775:       PetscCall(VecAXPBYPCZ(th->Xdot, -th->shift, th->shift, 0, th->X0, ts->vec_sol));
    776:       PetscCall(TSComputeIJacobianP(ts, th->stage_time, ts->vec_sol, th->Xdot, th->shift, ts->Jacp, PETSC_FALSE));
    777:       PetscCall(MatAXPY(MatDeltaFwdSensip, -1., ts->Jacp, SUBSET_NONZERO_PATTERN));
    778:     }
    779:   } else { /* 1-stage method */
    780:     th->shift = 0.0;
    781:     PetscCall(TSComputeIJacobian(ts, th->stage_time, th->X, th->Xdot, th->shift, J, Jpre, PETSC_FALSE));
    782:     PetscCall(MatMatMult(J, ts->mat_sensip, MAT_REUSE_MATRIX, PETSC_DETERMINE, &MatDeltaFwdSensip));
    783:     PetscCall(MatScale(MatDeltaFwdSensip, -1.));

    785:     /* Add the f_p forcing terms */
    786:     if (ts->Jacp) {
    787:       th->shift = previous_shift;
    788:       PetscCall(VecAXPBYPCZ(th->Xdot, -th->shift, th->shift, 0, th->X0, th->X));
    789:       PetscCall(TSComputeIJacobianP(ts, th->stage_time, th->X, th->Xdot, th->shift, ts->Jacp, PETSC_FALSE));
    790:       PetscCall(MatAXPY(MatDeltaFwdSensip, -1., ts->Jacp, SUBSET_NONZERO_PATTERN));
    791:     }
    792:   }

    794:   /* Build LHS */
    795:   th->shift = previous_shift; /* recover the previous shift used in TSStep_Theta() */
    796:   if (th->endpoint) {
    797:     PetscCall(TSComputeIJacobian(ts, th->stage_time, ts->vec_sol, th->Xdot, th->shift, J, Jpre, PETSC_FALSE));
    798:   } else {
    799:     PetscCall(TSComputeIJacobian(ts, th->stage_time, th->X, th->Xdot, th->shift, J, Jpre, PETSC_FALSE));
    800:   }
    801:   PetscCall(KSPSetOperators(ksp, J, Jpre));

    803:   /*
    804:     Evaluate the first stage of integral gradients with the 2-stage method:
    805:     drdu|t_n*S(t_n) + drdp|t_n
    806:     This is done before the linear solve because the sensitivity variable S(t_n) will be propagated to S(t_{n+1})
    807:   */
    808:   if (th->endpoint) { /* 2-stage method only */
    809:     if (quadts && quadts->mat_sensip) {
    810:       PetscCall(TSComputeRHSJacobian(quadts, th->ptime0, th->X0, quadJ, NULL));
    811:       PetscCall(TSComputeRHSJacobianP(quadts, th->ptime0, th->X0, quadJp));
    812:       PetscCall(MatTransposeMatMult(ts->mat_sensip, quadJ, MAT_REUSE_MATRIX, PETSC_DETERMINE, &th->MatIntegralSensipTemp));
    813:       PetscCall(MatAXPY(th->MatIntegralSensipTemp, 1, quadJp, SAME_NONZERO_PATTERN));
    814:       PetscCall(MatAXPY(quadts->mat_sensip, th->time_step0 * (1. - th->Theta), th->MatIntegralSensipTemp, SAME_NONZERO_PATTERN));
    815:     }
    816:   }

    818:   /* Solve the tangent linear equation for forward sensitivities to parameters */
    819:   for (ntlm = 0; ntlm < th->num_tlm; ntlm++) {
    820:     KSPConvergedReason kspreason;
    821:     PetscCall(MatDenseGetColumn(MatDeltaFwdSensip, ntlm, &barr));
    822:     PetscCall(VecPlaceArray(VecDeltaFwdSensipCol, barr));
    823:     if (th->endpoint) {
    824:       PetscCall(MatDenseGetColumn(ts->mat_sensip, ntlm, &xarr));
    825:       PetscCall(VecPlaceArray(ts->vec_sensip_col, xarr));
    826:       PetscCall(KSPSolve(ksp, VecDeltaFwdSensipCol, ts->vec_sensip_col));
    827:       PetscCall(VecResetArray(ts->vec_sensip_col));
    828:       PetscCall(MatDenseRestoreColumn(ts->mat_sensip, &xarr));
    829:     } else {
    830:       PetscCall(KSPSolve(ksp, VecDeltaFwdSensipCol, VecDeltaFwdSensipCol));
    831:     }
    832:     PetscCall(KSPGetConvergedReason(ksp, &kspreason));
    833:     if (kspreason < 0) {
    834:       ts->reason = TSFORWARD_DIVERGED_LINEAR_SOLVE;
    835:       PetscCall(PetscInfo(ts, "Step=%" PetscInt_FMT ", %" PetscInt_FMT "th tangent linear solve, linear solve fails, stopping tangent linear solve\n", ts->steps, ntlm));
    836:     }
    837:     PetscCall(VecResetArray(VecDeltaFwdSensipCol));
    838:     PetscCall(MatDenseRestoreColumn(MatDeltaFwdSensip, &barr));
    839:   }

    841:   /*
    842:     Evaluate the second stage of integral gradients with the 2-stage method:
    843:     drdu|t_{n+1}*S(t_{n+1}) + drdp|t_{n+1}
    844:   */
    845:   if (quadts && quadts->mat_sensip) {
    846:     if (!th->endpoint) {
    847:       PetscCall(MatAXPY(ts->mat_sensip, 1, MatDeltaFwdSensip, SAME_NONZERO_PATTERN)); /* stage sensitivity */
    848:       PetscCall(TSComputeRHSJacobian(quadts, th->stage_time, th->X, quadJ, NULL));
    849:       PetscCall(TSComputeRHSJacobianP(quadts, th->stage_time, th->X, quadJp));
    850:       PetscCall(MatTransposeMatMult(ts->mat_sensip, quadJ, MAT_REUSE_MATRIX, PETSC_DETERMINE, &th->MatIntegralSensipTemp));
    851:       PetscCall(MatAXPY(th->MatIntegralSensipTemp, 1, quadJp, SAME_NONZERO_PATTERN));
    852:       PetscCall(MatAXPY(quadts->mat_sensip, th->time_step0, th->MatIntegralSensipTemp, SAME_NONZERO_PATTERN));
    853:       PetscCall(MatAXPY(ts->mat_sensip, (1. - th->Theta) / th->Theta, MatDeltaFwdSensip, SAME_NONZERO_PATTERN));
    854:     } else {
    855:       PetscCall(TSComputeRHSJacobian(quadts, th->stage_time, ts->vec_sol, quadJ, NULL));
    856:       PetscCall(TSComputeRHSJacobianP(quadts, th->stage_time, ts->vec_sol, quadJp));
    857:       PetscCall(MatTransposeMatMult(ts->mat_sensip, quadJ, MAT_REUSE_MATRIX, PETSC_DETERMINE, &th->MatIntegralSensipTemp));
    858:       PetscCall(MatAXPY(th->MatIntegralSensipTemp, 1, quadJp, SAME_NONZERO_PATTERN));
    859:       PetscCall(MatAXPY(quadts->mat_sensip, th->time_step0 * th->Theta, th->MatIntegralSensipTemp, SAME_NONZERO_PATTERN));
    860:     }
    861:   } else {
    862:     if (!th->endpoint) PetscCall(MatAXPY(ts->mat_sensip, 1. / th->Theta, MatDeltaFwdSensip, SAME_NONZERO_PATTERN));
    863:   }
    864:   PetscFunctionReturn(PETSC_SUCCESS);
    865: }

    867: static PetscErrorCode TSForwardGetStages_Theta(TS ts, PetscInt *ns, Mat *stagesensip[])
    868: {
    869:   TS_Theta *th = (TS_Theta *)ts->data;

    871:   PetscFunctionBegin;
    872:   if (ns) {
    873:     if (!th->endpoint && th->Theta != 1.0) *ns = 1; /* midpoint form */
    874:     else *ns = 2;                                   /* endpoint form */
    875:   }
    876:   if (stagesensip) {
    877:     if (!th->endpoint && th->Theta != 1.0) {
    878:       th->MatFwdStages[0] = th->MatDeltaFwdSensip;
    879:     } else {
    880:       th->MatFwdStages[0] = th->MatFwdSensip0;
    881:       th->MatFwdStages[1] = ts->mat_sensip; /* stiffly accurate */
    882:     }
    883:     *stagesensip = th->MatFwdStages;
    884:   }
    885:   PetscFunctionReturn(PETSC_SUCCESS);
    886: }

    888: /*------------------------------------------------------------*/
    889: static PetscErrorCode TSReset_Theta(TS ts)
    890: {
    891:   TS_Theta *th = (TS_Theta *)ts->data;

    893:   PetscFunctionBegin;
    894:   PetscCall(VecDestroy(&th->X));
    895:   PetscCall(VecDestroy(&th->Xdot));
    896:   PetscCall(VecDestroy(&th->X0));
    897:   PetscCall(VecDestroy(&th->affine));

    899:   PetscCall(VecDestroy(&th->vec_sol_prev));
    900:   PetscCall(VecDestroy(&th->vec_lte_work));

    902:   PetscCall(VecDestroy(&th->VecCostIntegral0));
    903:   PetscFunctionReturn(PETSC_SUCCESS);
    904: }

    906: static PetscErrorCode TSAdjointReset_Theta(TS ts)
    907: {
    908:   TS_Theta *th = (TS_Theta *)ts->data;

    910:   PetscFunctionBegin;
    911:   PetscCall(VecDestroyVecs(ts->numcost, &th->VecsDeltaLam));
    912:   PetscCall(VecDestroyVecs(ts->numcost, &th->VecsDeltaMu));
    913:   PetscCall(VecDestroyVecs(ts->numcost, &th->VecsDeltaLam2));
    914:   PetscCall(VecDestroyVecs(ts->numcost, &th->VecsDeltaMu2));
    915:   PetscCall(VecDestroyVecs(ts->numcost, &th->VecsSensiTemp));
    916:   PetscCall(VecDestroyVecs(ts->numcost, &th->VecsSensi2Temp));
    917:   PetscFunctionReturn(PETSC_SUCCESS);
    918: }

    920: static PetscErrorCode TSDestroy_Theta(TS ts)
    921: {
    922:   PetscFunctionBegin;
    923:   PetscCall(TSReset_Theta(ts));
    924:   if (ts->dm) {
    925:     PetscCall(DMCoarsenHookRemove(ts->dm, DMCoarsenHook_TSTheta, DMRestrictHook_TSTheta, ts));
    926:     PetscCall(DMSubDomainHookRemove(ts->dm, DMSubDomainHook_TSTheta, DMSubDomainRestrictHook_TSTheta, ts));
    927:   }
    928:   PetscCall(PetscFree(ts->data));
    929:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaGetTheta_C", NULL));
    930:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaSetTheta_C", NULL));
    931:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaGetEndpoint_C", NULL));
    932:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaSetEndpoint_C", NULL));
    933:   PetscFunctionReturn(PETSC_SUCCESS);
    934: }

    936: /*
    937:   This defines the nonlinear equation that is to be solved with SNES
    938:   G(U) = F[t0+Theta*dt, U, (U-U0)*shift] = 0

    940:   Note that U here is the stage argument. This means that U = U_{n+1} only if endpoint = true,
    941:   otherwise U = theta U_{n+1} + (1 - theta) U0, which for the case of implicit midpoint is
    942:   U = (U_{n+1} + U0)/2
    943: */
    944: static PetscErrorCode SNESTSFormFunction_Theta(SNES snes, Vec x, Vec y, TS ts)
    945: {
    946:   TS_Theta *th = (TS_Theta *)ts->data;
    947:   Vec       X0, Xdot;
    948:   DM        dm, dmsave;
    949:   PetscReal shift = th->shift;

    951:   PetscFunctionBegin;
    952:   PetscCall(SNESGetDM(snes, &dm));
    953:   /* When using the endpoint variant, this is actually 1/Theta * Xdot */
    954:   PetscCall(TSThetaGetX0AndXdot(ts, dm, &X0, &Xdot));
    955:   if (x != X0) {
    956:     PetscCall(VecAXPBYPCZ(Xdot, -shift, shift, 0, X0, x));
    957:   } else {
    958:     PetscCall(VecZeroEntries(Xdot));
    959:   }
    960:   /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */
    961:   dmsave = ts->dm;
    962:   ts->dm = dm;
    963:   PetscCall(TSComputeIFunction(ts, th->stage_time, x, Xdot, y, PETSC_FALSE));
    964:   ts->dm = dmsave;
    965:   PetscCall(TSThetaRestoreX0AndXdot(ts, dm, &X0, &Xdot));
    966:   PetscFunctionReturn(PETSC_SUCCESS);
    967: }

    969: static PetscErrorCode SNESTSFormJacobian_Theta(SNES snes, Vec x, Mat A, Mat B, TS ts)
    970: {
    971:   TS_Theta *th = (TS_Theta *)ts->data;
    972:   Vec       Xdot;
    973:   DM        dm, dmsave;
    974:   PetscReal shift = th->shift;

    976:   PetscFunctionBegin;
    977:   PetscCall(SNESGetDM(snes, &dm));
    978:   /* Xdot has already been computed in SNESTSFormFunction_Theta (SNES guarantees this) */
    979:   PetscCall(TSThetaGetX0AndXdot(ts, dm, NULL, &Xdot));

    981:   dmsave = ts->dm;
    982:   ts->dm = dm;
    983:   PetscCall(TSComputeIJacobian(ts, th->stage_time, x, Xdot, shift, A, B, PETSC_FALSE));
    984:   ts->dm = dmsave;
    985:   PetscCall(TSThetaRestoreX0AndXdot(ts, dm, NULL, &Xdot));
    986:   PetscFunctionReturn(PETSC_SUCCESS);
    987: }

    989: static PetscErrorCode TSForwardSetUp_Theta(TS ts)
    990: {
    991:   TS_Theta *th     = (TS_Theta *)ts->data;
    992:   TS        quadts = ts->quadraturets;

    994:   PetscFunctionBegin;
    995:   /* combine sensitivities to parameters and sensitivities to initial values into one array */
    996:   th->num_tlm = ts->num_parameters;
    997:   PetscCall(MatDuplicate(ts->mat_sensip, MAT_DO_NOT_COPY_VALUES, &th->MatDeltaFwdSensip));
    998:   if (quadts && quadts->mat_sensip) {
    999:     PetscCall(MatDuplicate(quadts->mat_sensip, MAT_DO_NOT_COPY_VALUES, &th->MatIntegralSensipTemp));
    1000:     PetscCall(MatDuplicate(quadts->mat_sensip, MAT_DO_NOT_COPY_VALUES, &th->MatIntegralSensip0));
    1001:   }
    1002:   /* backup sensitivity results for roll-backs */
    1003:   PetscCall(MatDuplicate(ts->mat_sensip, MAT_DO_NOT_COPY_VALUES, &th->MatFwdSensip0));

    1005:   PetscCall(VecDuplicate(ts->vec_sol, &th->VecDeltaFwdSensipCol));
    1006:   PetscFunctionReturn(PETSC_SUCCESS);
    1007: }

    1009: static PetscErrorCode TSForwardReset_Theta(TS ts)
    1010: {
    1011:   TS_Theta *th     = (TS_Theta *)ts->data;
    1012:   TS        quadts = ts->quadraturets;

    1014:   PetscFunctionBegin;
    1015:   if (quadts && quadts->mat_sensip) {
    1016:     PetscCall(MatDestroy(&th->MatIntegralSensipTemp));
    1017:     PetscCall(MatDestroy(&th->MatIntegralSensip0));
    1018:   }
    1019:   PetscCall(VecDestroy(&th->VecDeltaFwdSensipCol));
    1020:   PetscCall(MatDestroy(&th->MatDeltaFwdSensip));
    1021:   PetscCall(MatDestroy(&th->MatFwdSensip0));
    1022:   PetscFunctionReturn(PETSC_SUCCESS);
    1023: }

    1025: static PetscErrorCode TSSetUp_Theta(TS ts)
    1026: {
    1027:   TS_Theta *th     = (TS_Theta *)ts->data;
    1028:   TS        quadts = ts->quadraturets;
    1029:   PetscBool match;

    1031:   PetscFunctionBegin;
    1032:   if (!th->VecCostIntegral0 && quadts && ts->costintegralfwd) { /* back up cost integral */
    1033:     PetscCall(VecDuplicate(quadts->vec_sol, &th->VecCostIntegral0));
    1034:   }
    1035:   if (!th->X) PetscCall(VecDuplicate(ts->vec_sol, &th->X));
    1036:   if (!th->Xdot) PetscCall(VecDuplicate(ts->vec_sol, &th->Xdot));
    1037:   if (!th->X0) PetscCall(VecDuplicate(ts->vec_sol, &th->X0));
    1038:   if (th->endpoint) PetscCall(VecDuplicate(ts->vec_sol, &th->affine));

    1040:   th->order = (th->Theta == 0.5) ? 2 : 1;
    1041:   th->shift = 1 / (th->Theta * ts->time_step);

    1043:   PetscCall(TSGetDM(ts, &ts->dm));
    1044:   PetscCall(DMCoarsenHookAdd(ts->dm, DMCoarsenHook_TSTheta, DMRestrictHook_TSTheta, ts));
    1045:   PetscCall(DMSubDomainHookAdd(ts->dm, DMSubDomainHook_TSTheta, DMSubDomainRestrictHook_TSTheta, ts));

    1047:   PetscCall(TSGetAdapt(ts, &ts->adapt));
    1048:   PetscCall(TSAdaptCandidatesClear(ts->adapt));
    1049:   PetscCall(PetscObjectTypeCompare((PetscObject)ts->adapt, TSADAPTNONE, &match));
    1050:   if (!match) {
    1051:     if (!th->vec_sol_prev) PetscCall(VecDuplicate(ts->vec_sol, &th->vec_sol_prev));
    1052:     if (!th->vec_lte_work) PetscCall(VecDuplicate(ts->vec_sol, &th->vec_lte_work));
    1053:   }
    1054:   PetscCall(TSGetSNES(ts, &ts->snes));

    1056:   ts->stifflyaccurate = (!th->endpoint && th->Theta != 1.0) ? PETSC_FALSE : PETSC_TRUE;
    1057:   PetscFunctionReturn(PETSC_SUCCESS);
    1058: }

    1060: /*------------------------------------------------------------*/

    1062: static PetscErrorCode TSAdjointSetUp_Theta(TS ts)
    1063: {
    1064:   TS_Theta *th = (TS_Theta *)ts->data;

    1066:   PetscFunctionBegin;
    1067:   PetscCall(VecDuplicateVecs(ts->vecs_sensi[0], ts->numcost, &th->VecsDeltaLam));
    1068:   PetscCall(VecDuplicateVecs(ts->vecs_sensi[0], ts->numcost, &th->VecsSensiTemp));
    1069:   if (ts->vecs_sensip) PetscCall(VecDuplicateVecs(ts->vecs_sensip[0], ts->numcost, &th->VecsDeltaMu));
    1070:   if (ts->vecs_sensi2) {
    1071:     PetscCall(VecDuplicateVecs(ts->vecs_sensi[0], ts->numcost, &th->VecsDeltaLam2));
    1072:     PetscCall(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:     PetscCall(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:   PetscFunctionReturn(PETSC_SUCCESS);
    1084: }

    1086: static PetscErrorCode TSSetFromOptions_Theta(TS ts, PetscOptionItems PetscOptionsObject)
    1087: {
    1088:   TS_Theta *th = (TS_Theta *)ts->data;

    1090:   PetscFunctionBegin;
    1091:   PetscOptionsHeadBegin(PetscOptionsObject, "Theta ODE solver options");
    1092:   {
    1093:     PetscCall(PetscOptionsReal("-ts_theta_theta", "Location of stage (0<Theta<=1)", "TSThetaSetTheta", th->Theta, &th->Theta, NULL));
    1094:     PetscCall(PetscOptionsBool("-ts_theta_endpoint", "Use the endpoint instead of midpoint form of the Theta method", "TSThetaSetEndpoint", th->endpoint, &th->endpoint, NULL));
    1095:     PetscCall(PetscOptionsBool("-ts_theta_initial_guess_extrapolate", "Extrapolate stage initial guess from previous solution (sometimes unstable)", "TSThetaSetExtrapolate", th->extrapolate, &th->extrapolate, NULL));
    1096:   }
    1097:   PetscOptionsHeadEnd();
    1098:   PetscFunctionReturn(PETSC_SUCCESS);
    1099: }

    1101: static PetscErrorCode TSView_Theta(TS ts, PetscViewer viewer)
    1102: {
    1103:   TS_Theta *th = (TS_Theta *)ts->data;
    1104:   PetscBool isascii;

    1106:   PetscFunctionBegin;
    1107:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
    1108:   if (isascii) {
    1109:     PetscCall(PetscViewerASCIIPrintf(viewer, "  Theta=%g\n", (double)th->Theta));
    1110:     PetscCall(PetscViewerASCIIPrintf(viewer, "  Extrapolation=%s\n", th->extrapolate ? "yes" : "no"));
    1111:   }
    1112:   PetscFunctionReturn(PETSC_SUCCESS);
    1113: }

    1115: static PetscErrorCode TSThetaGetTheta_Theta(TS ts, PetscReal *theta)
    1116: {
    1117:   TS_Theta *th = (TS_Theta *)ts->data;

    1119:   PetscFunctionBegin;
    1120:   *theta = th->Theta;
    1121:   PetscFunctionReturn(PETSC_SUCCESS);
    1122: }

    1124: static PetscErrorCode TSThetaSetTheta_Theta(TS ts, PetscReal theta)
    1125: {
    1126:   TS_Theta *th = (TS_Theta *)ts->data;

    1128:   PetscFunctionBegin;
    1129:   PetscCheck(theta > 0 && theta <= 1, PetscObjectComm((PetscObject)ts), PETSC_ERR_ARG_OUTOFRANGE, "Theta %g not in range (0,1]", (double)theta);
    1130:   th->Theta = theta;
    1131:   th->order = (th->Theta == 0.5) ? 2 : 1;
    1132:   PetscFunctionReturn(PETSC_SUCCESS);
    1133: }

    1135: static PetscErrorCode TSThetaGetEndpoint_Theta(TS ts, PetscBool *endpoint)
    1136: {
    1137:   TS_Theta *th = (TS_Theta *)ts->data;

    1139:   PetscFunctionBegin;
    1140:   *endpoint = th->endpoint;
    1141:   PetscFunctionReturn(PETSC_SUCCESS);
    1142: }

    1144: static PetscErrorCode TSThetaSetEndpoint_Theta(TS ts, PetscBool flg)
    1145: {
    1146:   TS_Theta *th = (TS_Theta *)ts->data;

    1148:   PetscFunctionBegin;
    1149:   th->endpoint = flg;
    1150:   PetscFunctionReturn(PETSC_SUCCESS);
    1151: }

    1153: #if defined(PETSC_HAVE_COMPLEX)
    1154: static PetscErrorCode TSComputeLinearStability_Theta(TS ts, PetscReal xr, PetscReal xi, PetscReal *yr, PetscReal *yi)
    1155: {
    1156:   PetscComplex z  = xr + xi * PETSC_i, f;
    1157:   TS_Theta    *th = (TS_Theta *)ts->data;

    1159:   PetscFunctionBegin;
    1160:   f   = (1.0 + (1.0 - th->Theta) * z) / (1.0 - th->Theta * z);
    1161:   *yr = PetscRealPartComplex(f);
    1162:   *yi = PetscImaginaryPartComplex(f);
    1163:   PetscFunctionReturn(PETSC_SUCCESS);
    1164: }
    1165: #endif

    1167: static PetscErrorCode TSGetStages_Theta(TS ts, PetscInt *ns, Vec *Y[])
    1168: {
    1169:   TS_Theta *th = (TS_Theta *)ts->data;

    1171:   PetscFunctionBegin;
    1172:   if (ns) {
    1173:     if (!th->endpoint && th->Theta != 1.0) *ns = 1; /* midpoint form */
    1174:     else *ns = 2;                                   /* endpoint form */
    1175:   }
    1176:   if (Y) {
    1177:     if (!th->endpoint && th->Theta != 1.0) {
    1178:       th->Stages[0] = th->X;
    1179:     } else {
    1180:       th->Stages[0] = th->X0;
    1181:       th->Stages[1] = ts->vec_sol; /* stiffly accurate */
    1182:     }
    1183:     *Y = th->Stages;
    1184:   }
    1185:   PetscFunctionReturn(PETSC_SUCCESS);
    1186: }

    1188: /* ------------------------------------------------------------ */
    1189: /*MC
    1190:       TSTHETA - DAE solver using the implicit Theta method

    1192:    Level: beginner

    1194:    Options Database Keys:
    1195: +  -ts_theta_theta <Theta> - Location of stage (0<Theta<=1)
    1196: .  -ts_theta_endpoint <flag> - Use the endpoint (like Crank-Nicholson) instead of midpoint form of the Theta method
    1197: -  -ts_theta_initial_guess_extrapolate <flg> - Extrapolate stage initial guess from previous solution (sometimes unstable)

    1199:    Notes:
    1200: .vb
    1201:   -ts_type theta -ts_theta_theta 1.0 corresponds to backward Euler (TSBEULER)
    1202:   -ts_type theta -ts_theta_theta 0.5 corresponds to the implicit midpoint rule
    1203:   -ts_type theta -ts_theta_theta 0.5 -ts_theta_endpoint corresponds to Crank-Nicholson (TSCN)
    1204: .ve

    1206:    The endpoint variant of the Theta method and backward Euler can be applied to DAE. The midpoint variant is not suitable for DAEs because it is not stiffly accurate.

    1208:    The midpoint variant is cast as a 1-stage implicit Runge-Kutta method.

    1210: .vb
    1211:   Theta | Theta
    1212:   -------------
    1213:         |  1
    1214: .ve

    1216:    For the default Theta=0.5, this is also known as the implicit midpoint rule.

    1218:    When the endpoint variant is chosen, the method becomes a 2-stage method with first stage explicit:

    1220: .vb
    1221:   0 | 0         0
    1222:   1 | 1-Theta   Theta
    1223:   -------------------
    1224:     | 1-Theta   Theta
    1225: .ve

    1227:    For the default Theta=0.5, this is the trapezoid rule (also known as Crank-Nicolson, see TSCN).

    1229:    To apply a diagonally implicit RK method to DAE, the stage formula
    1230: .vb
    1231:   Y_i = X + h sum_j a_ij Y'_j
    1232: .ve
    1233:    is interpreted as a formula for Y'_i in terms of Y_i and known values (Y'_j, j<i)

    1235: .seealso: [](ch_ts), `TSCreate()`, `TS`, `TSSetType()`, `TSCN`, `TSBEULER`, `TSThetaSetTheta()`, `TSThetaSetEndpoint()`
    1236: M*/
    1237: PETSC_EXTERN PetscErrorCode TSCreate_Theta(TS ts)
    1238: {
    1239:   TS_Theta *th;

    1241:   PetscFunctionBegin;
    1242:   ts->ops->reset          = TSReset_Theta;
    1243:   ts->ops->adjointreset   = TSAdjointReset_Theta;
    1244:   ts->ops->destroy        = TSDestroy_Theta;
    1245:   ts->ops->view           = TSView_Theta;
    1246:   ts->ops->setup          = TSSetUp_Theta;
    1247:   ts->ops->adjointsetup   = TSAdjointSetUp_Theta;
    1248:   ts->ops->adjointreset   = TSAdjointReset_Theta;
    1249:   ts->ops->step           = TSStep_Theta;
    1250:   ts->ops->interpolate    = TSInterpolate_Theta;
    1251:   ts->ops->evaluatewlte   = TSEvaluateWLTE_Theta;
    1252:   ts->ops->rollback       = TSRollBack_Theta;
    1253:   ts->ops->resizeregister = TSResizeRegister_Theta;
    1254:   ts->ops->setfromoptions = TSSetFromOptions_Theta;
    1255:   ts->ops->snesfunction   = SNESTSFormFunction_Theta;
    1256:   ts->ops->snesjacobian   = SNESTSFormJacobian_Theta;
    1257: #if defined(PETSC_HAVE_COMPLEX)
    1258:   ts->ops->linearstability = TSComputeLinearStability_Theta;
    1259: #endif
    1260:   ts->ops->getstages       = TSGetStages_Theta;
    1261:   ts->ops->adjointstep     = TSAdjointStep_Theta;
    1262:   ts->ops->adjointintegral = TSAdjointCostIntegral_Theta;
    1263:   ts->ops->forwardintegral = TSForwardCostIntegral_Theta;
    1264:   ts->default_adapt_type   = TSADAPTNONE;

    1266:   ts->ops->forwardsetup     = TSForwardSetUp_Theta;
    1267:   ts->ops->forwardreset     = TSForwardReset_Theta;
    1268:   ts->ops->forwardstep      = TSForwardStep_Theta;
    1269:   ts->ops->forwardgetstages = TSForwardGetStages_Theta;

    1271:   ts->usessnes = PETSC_TRUE;

    1273:   PetscCall(PetscNew(&th));
    1274:   ts->data = (void *)th;

    1276:   th->VecsDeltaLam   = NULL;
    1277:   th->VecsDeltaMu    = NULL;
    1278:   th->VecsSensiTemp  = NULL;
    1279:   th->VecsSensi2Temp = NULL;

    1281:   th->extrapolate = PETSC_FALSE;
    1282:   th->Theta       = 0.5;
    1283:   th->order       = 2;
    1284:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaGetTheta_C", TSThetaGetTheta_Theta));
    1285:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaSetTheta_C", TSThetaSetTheta_Theta));
    1286:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaGetEndpoint_C", TSThetaGetEndpoint_Theta));
    1287:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaSetEndpoint_C", TSThetaSetEndpoint_Theta));
    1288:   PetscFunctionReturn(PETSC_SUCCESS);
    1289: }

    1291: /*@
    1292:   TSThetaGetTheta - Get the abscissa of the stage in (0,1] for `TSTHETA`

    1294:   Not Collective

    1296:   Input Parameter:
    1297: . ts - timestepping context

    1299:   Output Parameter:
    1300: . theta - stage abscissa

    1302:   Level: advanced

    1304:   Note:
    1305:   Use of this function is normally only required to hack `TSTHETA` to use a modified integration scheme.

    1307: .seealso: [](ch_ts), `TSThetaSetTheta()`, `TSTHETA`
    1308: @*/
    1309: PetscErrorCode TSThetaGetTheta(TS ts, PetscReal *theta)
    1310: {
    1311:   PetscFunctionBegin;
    1313:   PetscAssertPointer(theta, 2);
    1314:   PetscUseMethod(ts, "TSThetaGetTheta_C", (TS, PetscReal *), (ts, theta));
    1315:   PetscFunctionReturn(PETSC_SUCCESS);
    1316: }

    1318: /*@
    1319:   TSThetaSetTheta - Set the abscissa of the stage in (0,1]  for `TSTHETA`

    1321:   Not Collective

    1323:   Input Parameters:
    1324: + ts    - timestepping context
    1325: - theta - stage abscissa

    1327:   Options Database Key:
    1328: . -ts_theta_theta <theta> - set theta

    1330:   Level: intermediate

    1332: .seealso: [](ch_ts), `TSThetaGetTheta()`, `TSTHETA`, `TSCN`
    1333: @*/
    1334: PetscErrorCode TSThetaSetTheta(TS ts, PetscReal theta)
    1335: {
    1336:   PetscFunctionBegin;
    1338:   PetscTryMethod(ts, "TSThetaSetTheta_C", (TS, PetscReal), (ts, theta));
    1339:   PetscFunctionReturn(PETSC_SUCCESS);
    1340: }

    1342: /*@
    1343:   TSThetaGetEndpoint - Gets whether to use the endpoint variant of the method (e.g. trapezoid/Crank-Nicolson instead of midpoint rule) for `TSTHETA`

    1345:   Not Collective

    1347:   Input Parameter:
    1348: . ts - timestepping context

    1350:   Output Parameter:
    1351: . endpoint - `PETSC_TRUE` when using the endpoint variant

    1353:   Level: advanced

    1355: .seealso: [](ch_ts), `TSThetaSetEndpoint()`, `TSTHETA`, `TSCN`
    1356: @*/
    1357: PetscErrorCode TSThetaGetEndpoint(TS ts, PetscBool *endpoint)
    1358: {
    1359:   PetscFunctionBegin;
    1361:   PetscAssertPointer(endpoint, 2);
    1362:   PetscUseMethod(ts, "TSThetaGetEndpoint_C", (TS, PetscBool *), (ts, endpoint));
    1363:   PetscFunctionReturn(PETSC_SUCCESS);
    1364: }

    1366: /*@
    1367:   TSThetaSetEndpoint - Sets whether to use the endpoint variant of the method (e.g. trapezoid/Crank-Nicolson instead of midpoint rule) for `TSTHETA`

    1369:   Not Collective

    1371:   Input Parameters:
    1372: + ts  - timestepping context
    1373: - flg - `PETSC_TRUE` to use the endpoint variant

    1375:   Options Database Key:
    1376: . -ts_theta_endpoint <flg> - use the endpoint variant

    1378:   Level: intermediate

    1380: .seealso: [](ch_ts), `TSTHETA`, `TSCN`
    1381: @*/
    1382: PetscErrorCode TSThetaSetEndpoint(TS ts, PetscBool flg)
    1383: {
    1384:   PetscFunctionBegin;
    1386:   PetscTryMethod(ts, "TSThetaSetEndpoint_C", (TS, PetscBool), (ts, flg));
    1387:   PetscFunctionReturn(PETSC_SUCCESS);
    1388: }

    1390: /*
    1391:  * TSBEULER and TSCN are straightforward specializations of TSTHETA.
    1392:  * The creation functions for these specializations are below.
    1393:  */

    1395: static PetscErrorCode TSSetUp_BEuler(TS ts)
    1396: {
    1397:   TS_Theta *th = (TS_Theta *)ts->data;

    1399:   PetscFunctionBegin;
    1400:   PetscCheck(th->Theta == 1.0, PetscObjectComm((PetscObject)ts), PETSC_ERR_OPT_OVERWRITE, "Can not change the default value (1) of theta when using backward Euler");
    1401:   PetscCheck(!th->endpoint, PetscObjectComm((PetscObject)ts), PETSC_ERR_OPT_OVERWRITE, "Can not change to the endpoint form of the Theta methods when using backward Euler");
    1402:   PetscCall(TSSetUp_Theta(ts));
    1403:   PetscFunctionReturn(PETSC_SUCCESS);
    1404: }

    1406: static PetscErrorCode TSView_BEuler(TS ts, PetscViewer viewer)
    1407: {
    1408:   PetscFunctionBegin;
    1409:   PetscFunctionReturn(PETSC_SUCCESS);
    1410: }

    1412: /*MC
    1413:       TSBEULER - ODE solver using the implicit backward Euler method

    1415:   Level: beginner

    1417:   Note:
    1418:   `TSBEULER` is equivalent to `TSTHETA` with Theta=1.0 or `-ts_type theta -ts_theta_theta 1.0`

    1420: .seealso: [](ch_ts), `TSCreate()`, `TS`, `TSSetType()`, `TSEULER`, `TSCN`, `TSTHETA`
    1421: M*/
    1422: PETSC_EXTERN PetscErrorCode TSCreate_BEuler(TS ts)
    1423: {
    1424:   PetscFunctionBegin;
    1425:   PetscCall(TSCreate_Theta(ts));
    1426:   PetscCall(TSThetaSetTheta(ts, 1.0));
    1427:   PetscCall(TSThetaSetEndpoint(ts, PETSC_FALSE));
    1428:   ts->ops->setup = TSSetUp_BEuler;
    1429:   ts->ops->view  = TSView_BEuler;
    1430:   PetscFunctionReturn(PETSC_SUCCESS);
    1431: }

    1433: static PetscErrorCode TSSetUp_CN(TS ts)
    1434: {
    1435:   TS_Theta *th = (TS_Theta *)ts->data;

    1437:   PetscFunctionBegin;
    1438:   PetscCheck(th->Theta == 0.5, PetscObjectComm((PetscObject)ts), PETSC_ERR_OPT_OVERWRITE, "Can not change the default value (0.5) of theta when using Crank-Nicolson");
    1439:   PetscCheck(th->endpoint, PetscObjectComm((PetscObject)ts), PETSC_ERR_OPT_OVERWRITE, "Can not change to the midpoint form of the Theta methods when using Crank-Nicolson");
    1440:   PetscCall(TSSetUp_Theta(ts));
    1441:   PetscFunctionReturn(PETSC_SUCCESS);
    1442: }

    1444: static PetscErrorCode TSView_CN(TS ts, PetscViewer viewer)
    1445: {
    1446:   PetscFunctionBegin;
    1447:   PetscFunctionReturn(PETSC_SUCCESS);
    1448: }

    1450: /*MC
    1451:       TSCN - ODE solver using the implicit Crank-Nicolson method.

    1453:   Level: beginner

    1455:   Notes:
    1456:   `TSCN` is equivalent to `TSTHETA` with Theta=0.5 and the "endpoint" option set. I.e.
    1457: .vb
    1458:   -ts_type theta
    1459:   -ts_theta_theta 0.5
    1460:   -ts_theta_endpoint
    1461: .ve

    1463: .seealso: [](ch_ts), `TSCreate()`, `TS`, `TSSetType()`, `TSBEULER`, `TSTHETA`, `TSType`,
    1464: M*/
    1465: PETSC_EXTERN PetscErrorCode TSCreate_CN(TS ts)
    1466: {
    1467:   PetscFunctionBegin;
    1468:   PetscCall(TSCreate_Theta(ts));
    1469:   PetscCall(TSThetaSetTheta(ts, 0.5));
    1470:   PetscCall(TSThetaSetEndpoint(ts, PETSC_TRUE));
    1471:   ts->ops->setup = TSSetUp_CN;
    1472:   ts->ops->view  = TSView_CN;
    1473:   PetscFunctionReturn(PETSC_SUCCESS);
    1474: }