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, PetscCtx ctx)
     74: {
     75:   PetscFunctionBegin;
     76:   PetscFunctionReturn(PETSC_SUCCESS);
     77: }

     79: static PetscErrorCode DMRestrictHook_TSTheta(DM fine, Mat restrct, Vec rscale, Mat inject, DM coarse, PetscCtx 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, PetscCtx ctx)
     97: {
     98:   PetscFunctionBegin;
     99:   PetscFunctionReturn(PETSC_SUCCESS);
    100: }

    102: static PetscErrorCode DMSubDomainRestrictHook_TSTheta(DM dm, VecScatter gscat, VecScatter lscat, DM subdm, PetscCtx 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) PetscCall(VecCopy(th->X, ts->vec_sol));
    229:     else {
    230:       PetscCall(VecAXPBYPCZ(th->Xdot, -th->shift, th->shift, 0, th->X0, th->X)); /* th->Xdot is needed by TSInterpolate_Theta */
    231:       if (th->Theta == 1.0) PetscCall(VecCopy(th->X, ts->vec_sol));              /* BEULER, stage already checked */
    232:       else {
    233:         PetscCall(VecAXPY(ts->vec_sol, ts->time_step, th->Xdot));
    234:         PetscCall(TSAdaptCheckStage(ts->adapt, ts, ts->ptime + ts->time_step, ts->vec_sol, &stageok));
    235:         if (!stageok) {
    236:           PetscCall(VecCopy(th->X0, ts->vec_sol));
    237:           goto reject_step;
    238:         }
    239:       }
    240:     }

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

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

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

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

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

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

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

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

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

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

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

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

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

    407:   if (ts->vecs_sensi2) {
    408:     PetscCall(VecResetArray(ts->vec_sensip_col));
    409:     PetscCall(MatDenseRestoreColumn(ts->mat_sensip, &xarr));
    410:   }
    411:   th->status = TS_STEP_COMPLETE;
    412:   PetscFunctionReturn(PETSC_SUCCESS);
    413: }

    415: static PetscErrorCode TSAdjointStep_Theta(TS ts)
    416: {
    417:   TS_Theta    *th           = (TS_Theta *)ts->data;
    418:   TS           quadts       = ts->quadraturets;
    419:   Vec         *VecsDeltaLam = th->VecsDeltaLam, *VecsDeltaMu = th->VecsDeltaMu, *VecsSensiTemp = th->VecsSensiTemp;
    420:   Vec         *VecsDeltaLam2 = th->VecsDeltaLam2, *VecsDeltaMu2 = th->VecsDeltaMu2, *VecsSensi2Temp = th->VecsSensi2Temp;
    421:   PetscInt     nadj;
    422:   Mat          J, Jpre, quadJ = NULL, quadJp = NULL;
    423:   KSP          ksp;
    424:   PetscScalar *xarr;
    425:   PetscReal    adjoint_time_step;
    426:   PetscReal    adjoint_ptime; /* end time of the adjoint time step (ts->ptime is the start time, usually ts->ptime is larger than adjoint_ptime) */

    428:   PetscFunctionBegin;
    429:   if (th->Theta == 1.) {
    430:     PetscCall(TSAdjointStepBEuler_Private(ts));
    431:     PetscFunctionReturn(PETSC_SUCCESS);
    432:   }
    433:   th->status = TS_STEP_INCOMPLETE;
    434:   PetscCall(SNESGetKSP(ts->snes, &ksp));
    435:   PetscCall(TSGetIJacobian(ts, &J, &Jpre, NULL, NULL));
    436:   if (quadts) {
    437:     PetscCall(TSGetRHSJacobian(quadts, &quadJ, NULL, NULL, NULL));
    438:     PetscCall(TSGetRHSJacobianP(quadts, &quadJp, NULL, NULL));
    439:   }
    440:   /* If endpoint=1, th->ptime and th->X0 will be used; if endpoint=0, th->stage_time and th->X will be used. */
    441:   th->stage_time    = th->endpoint ? ts->ptime : (ts->ptime + (1. - th->Theta) * ts->time_step);
    442:   adjoint_ptime     = ts->ptime + ts->time_step;
    443:   adjoint_time_step = -ts->time_step; /* always positive since time_step is negative */

    445:   if (!th->endpoint) {
    446:     /* recover th->X0 using vec_sol and the stage value th->X */
    447:     PetscCall(VecAXPBYPCZ(th->X0, 1.0 / (1.0 - th->Theta), th->Theta / (th->Theta - 1.0), 0, th->X, ts->vec_sol));
    448:   }

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

    457:   for (nadj = 0; nadj < ts->numcost; nadj++) {
    458:     PetscCall(VecCopy(ts->vecs_sensi[nadj], VecsSensiTemp[nadj]));
    459:     PetscCall(VecScale(VecsSensiTemp[nadj], 1. / (th->Theta * adjoint_time_step)));
    460:     if (quadJ) {
    461:       PetscCall(MatDenseGetColumn(quadJ, nadj, &xarr));
    462:       PetscCall(VecPlaceArray(ts->vec_drdu_col, xarr));
    463:       PetscCall(VecAXPY(VecsSensiTemp[nadj], 1., ts->vec_drdu_col));
    464:       PetscCall(VecResetArray(ts->vec_drdu_col));
    465:       PetscCall(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:     PetscCall(TSComputeSNESJacobian(ts, ts->vec_sol, J, Jpre));
    473:   } else {
    474:     PetscCall(TSComputeSNESJacobian(ts, th->X, J, Jpre));
    475:   }
    476:   PetscCall(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:     PetscCall(KSPSolveTranspose(ksp, VecsSensiTemp[nadj], VecsDeltaLam[nadj]));
    482:     PetscCall(KSPGetConvergedReason(ksp, &kspreason));
    483:     if (kspreason < 0) {
    484:       ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
    485:       PetscCall(PetscInfo(ts, "Step=%" PetscInt_FMT ", %" PetscInt_FMT "th 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:     PetscCheck(th->endpoint, PetscObjectComm((PetscObject)ts), PETSC_ERR_SUP, "Operation not implemented in TS_Theta");
    492:     /* Get w1 at t_{n+1} from TLM matrix */
    493:     PetscCall(MatDenseGetColumn(ts->mat_sensip, 0, &xarr));
    494:     PetscCall(VecPlaceArray(ts->vec_sensip_col, xarr));
    495:     /* lambda_s^T F_UU w_1 */
    496:     PetscCall(TSComputeIHessianProductFunctionUU(ts, th->stage_time, ts->vec_sol, VecsDeltaLam, ts->vec_sensip_col, ts->vecs_fuu));
    497:     PetscCall(VecResetArray(ts->vec_sensip_col));
    498:     PetscCall(MatDenseRestoreColumn(ts->mat_sensip, &xarr));
    499:     /* lambda_s^T F_UP w_2 */
    500:     PetscCall(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:       PetscCall(VecCopy(ts->vecs_sensi2[nadj], VecsSensi2Temp[nadj]));
    503:       PetscCall(VecScale(VecsSensi2Temp[nadj], th->shift));
    504:       PetscCall(VecAXPY(VecsSensi2Temp[nadj], -1., ts->vecs_fuu[nadj]));
    505:       if (ts->vecs_fup) PetscCall(VecAXPY(VecsSensi2Temp[nadj], -1., ts->vecs_fup[nadj]));
    506:     }
    507:     /* Solve stage equation LHS X = RHS for second-order adjoint */
    508:     for (nadj = 0; nadj < ts->numcost; nadj++) {
    509:       KSPConvergedReason kspreason;
    510:       PetscCall(KSPSolveTranspose(ksp, VecsSensi2Temp[nadj], VecsDeltaLam2[nadj]));
    511:       PetscCall(KSPGetConvergedReason(ksp, &kspreason));
    512:       if (kspreason < 0) {
    513:         ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
    514:         PetscCall(PetscInfo(ts, "Step=%" PetscInt_FMT ", %" PetscInt_FMT "th cost function, transposed linear solve fails, stopping 2nd-order adjoint solve\n", ts->steps, nadj));
    515:       }
    516:     }
    517:   }

    519:   /* Update sensitivities, and evaluate integrals if there is any */
    520:   if (th->endpoint) { /* two-stage Theta methods with th->Theta!=1, th->Theta==1 leads to BEuler */
    521:     th->shift      = 1. / ((th->Theta - 1.) * adjoint_time_step);
    522:     th->stage_time = adjoint_ptime;
    523:     PetscCall(TSComputeSNESJacobian(ts, th->X0, J, Jpre));
    524:     PetscCall(KSPSetOperators(ksp, J, Jpre));
    525:     /* R_U at t_n */
    526:     if (quadts) PetscCall(TSComputeRHSJacobian(quadts, adjoint_ptime, th->X0, quadJ, NULL));
    527:     for (nadj = 0; nadj < ts->numcost; nadj++) {
    528:       PetscCall(MatMultTranspose(J, VecsDeltaLam[nadj], ts->vecs_sensi[nadj]));
    529:       if (quadJ) {
    530:         PetscCall(MatDenseGetColumn(quadJ, nadj, &xarr));
    531:         PetscCall(VecPlaceArray(ts->vec_drdu_col, xarr));
    532:         PetscCall(VecAXPY(ts->vecs_sensi[nadj], -1., ts->vec_drdu_col));
    533:         PetscCall(VecResetArray(ts->vec_drdu_col));
    534:         PetscCall(MatDenseRestoreColumn(quadJ, &xarr));
    535:       }
    536:       PetscCall(VecScale(ts->vecs_sensi[nadj], 1. / th->shift));
    537:     }

    539:     /* Second-order adjoint */
    540:     if (ts->vecs_sensi2) { /* U_n */
    541:       /* Get w1 at t_n from TLM matrix */
    542:       PetscCall(MatDenseGetColumn(th->MatFwdSensip0, 0, &xarr));
    543:       PetscCall(VecPlaceArray(ts->vec_sensip_col, xarr));
    544:       /* lambda_s^T F_UU w_1 */
    545:       PetscCall(TSComputeIHessianProductFunctionUU(ts, adjoint_ptime, th->X0, VecsDeltaLam, ts->vec_sensip_col, ts->vecs_fuu));
    546:       PetscCall(VecResetArray(ts->vec_sensip_col));
    547:       PetscCall(MatDenseRestoreColumn(th->MatFwdSensip0, &xarr));
    548:       /* lambda_s^T F_UU w_2 */
    549:       PetscCall(TSComputeIHessianProductFunctionUP(ts, adjoint_ptime, th->X0, VecsDeltaLam, ts->vec_dir, ts->vecs_fup));
    550:       for (nadj = 0; nadj < ts->numcost; nadj++) {
    551:         /* 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  */
    552:         PetscCall(MatMultTranspose(J, VecsDeltaLam2[nadj], ts->vecs_sensi2[nadj]));
    553:         PetscCall(VecAXPY(ts->vecs_sensi2[nadj], 1., ts->vecs_fuu[nadj]));
    554:         if (ts->vecs_fup) PetscCall(VecAXPY(ts->vecs_sensi2[nadj], 1., ts->vecs_fup[nadj]));
    555:         PetscCall(VecScale(ts->vecs_sensi2[nadj], 1. / th->shift));
    556:       }
    557:     }

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

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

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

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

    667:   th->status = TS_STEP_COMPLETE;
    668:   PetscFunctionReturn(PETSC_SUCCESS);
    669: }

    671: static PetscErrorCode TSInterpolate_Theta(TS ts, PetscReal t, Vec X)
    672: {
    673:   TS_Theta *th = (TS_Theta *)ts->data;
    674:   PetscReal dt = t - ts->ptime;

    676:   PetscFunctionBegin;
    677:   PetscCall(VecCopy(ts->vec_sol, th->X));
    678:   if (th->endpoint) dt *= th->Theta;
    679:   PetscCall(VecWAXPY(X, dt, th->Xdot, th->X));
    680:   PetscFunctionReturn(PETSC_SUCCESS);
    681: }

    683: static PetscErrorCode TSEvaluateWLTE_Theta(TS ts, NormType wnormtype, PetscInt *order, PetscReal *wlte)
    684: {
    685:   TS_Theta *th = (TS_Theta *)ts->data;
    686:   Vec       X  = ts->vec_sol;      /* X = solution */
    687:   Vec       Y  = th->vec_lte_work; /* Y = X + LTE  */
    688:   PetscReal wltea, wlter;

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

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

    721: static PetscErrorCode TSRollBack_Theta(TS ts)
    722: {
    723:   TS_Theta *th     = (TS_Theta *)ts->data;
    724:   TS        quadts = ts->quadraturets;

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

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

    746:   PetscFunctionBegin;
    747:   previous_shift = th->shift;
    748:   PetscCall(MatCopy(ts->mat_sensip, th->MatFwdSensip0, SAME_NONZERO_PATTERN));

    750:   if (quadts && quadts->mat_sensip) PetscCall(MatCopy(quadts->mat_sensip, th->MatIntegralSensip0, SAME_NONZERO_PATTERN));
    751:   PetscCall(SNESGetKSP(ts->snes, &ksp));
    752:   PetscCall(TSGetIJacobian(ts, &J, &Jpre, NULL, NULL));
    753:   if (quadts) {
    754:     PetscCall(TSGetRHSJacobian(quadts, &quadJ, NULL, NULL, NULL));
    755:     PetscCall(TSGetRHSJacobianP(quadts, &quadJp, NULL, NULL));
    756:   }

    758:   /* Build RHS */
    759:   if (th->endpoint) { /* 2-stage method*/
    760:     th->shift = 1. / ((th->Theta - 1.) * th->time_step0);
    761:     PetscCall(TSComputeIJacobian(ts, th->ptime0, th->X0, th->Xdot, th->shift, J, Jpre, PETSC_FALSE));
    762:     PetscCall(MatMatMult(J, ts->mat_sensip, MAT_REUSE_MATRIX, PETSC_DETERMINE, &MatDeltaFwdSensip));
    763:     PetscCall(MatScale(MatDeltaFwdSensip, (th->Theta - 1.) / th->Theta));

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

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

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

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

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

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

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

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

    884: static PetscErrorCode TSReset_Theta(TS ts)
    885: {
    886:   TS_Theta *th = (TS_Theta *)ts->data;

    888:   PetscFunctionBegin;
    889:   PetscCall(VecDestroy(&th->X));
    890:   PetscCall(VecDestroy(&th->Xdot));
    891:   PetscCall(VecDestroy(&th->X0));
    892:   PetscCall(VecDestroy(&th->affine));

    894:   PetscCall(VecDestroy(&th->vec_sol_prev));
    895:   PetscCall(VecDestroy(&th->vec_lte_work));

    897:   PetscCall(VecDestroy(&th->VecCostIntegral0));
    898:   PetscFunctionReturn(PETSC_SUCCESS);
    899: }

    901: static PetscErrorCode TSAdjointReset_Theta(TS ts)
    902: {
    903:   TS_Theta *th = (TS_Theta *)ts->data;

    905:   PetscFunctionBegin;
    906:   PetscCall(VecDestroyVecs(ts->numcost, &th->VecsDeltaLam));
    907:   PetscCall(VecDestroyVecs(ts->numcost, &th->VecsDeltaMu));
    908:   PetscCall(VecDestroyVecs(ts->numcost, &th->VecsDeltaLam2));
    909:   PetscCall(VecDestroyVecs(ts->numcost, &th->VecsDeltaMu2));
    910:   PetscCall(VecDestroyVecs(ts->numcost, &th->VecsSensiTemp));
    911:   PetscCall(VecDestroyVecs(ts->numcost, &th->VecsSensi2Temp));
    912:   PetscFunctionReturn(PETSC_SUCCESS);
    913: }

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

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

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

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

    961: static PetscErrorCode SNESTSFormJacobian_Theta(SNES snes, Vec x, Mat A, Mat B, TS ts)
    962: {
    963:   TS_Theta *th = (TS_Theta *)ts->data;
    964:   Vec       Xdot;
    965:   DM        dm, dmsave;
    966:   PetscReal shift = th->shift;

    968:   PetscFunctionBegin;
    969:   PetscCall(SNESGetDM(snes, &dm));
    970:   /* Xdot has already been computed in SNESTSFormFunction_Theta (SNES guarantees this) */
    971:   PetscCall(TSThetaGetX0AndXdot(ts, dm, NULL, &Xdot));

    973:   dmsave = ts->dm;
    974:   ts->dm = dm;
    975:   PetscCall(TSComputeIJacobian(ts, th->stage_time, x, Xdot, shift, A, B, PETSC_FALSE));
    976:   ts->dm = dmsave;
    977:   PetscCall(TSThetaRestoreX0AndXdot(ts, dm, NULL, &Xdot));
    978:   PetscFunctionReturn(PETSC_SUCCESS);
    979: }

    981: static PetscErrorCode TSForwardSetUp_Theta(TS ts)
    982: {
    983:   TS_Theta *th     = (TS_Theta *)ts->data;
    984:   TS        quadts = ts->quadraturets;

    986:   PetscFunctionBegin;
    987:   /* combine sensitivities to parameters and sensitivities to initial values into one array */
    988:   th->num_tlm = ts->num_parameters;
    989:   PetscCall(MatDuplicate(ts->mat_sensip, MAT_DO_NOT_COPY_VALUES, &th->MatDeltaFwdSensip));
    990:   if (quadts && quadts->mat_sensip) {
    991:     PetscCall(MatDuplicate(quadts->mat_sensip, MAT_DO_NOT_COPY_VALUES, &th->MatIntegralSensipTemp));
    992:     PetscCall(MatDuplicate(quadts->mat_sensip, MAT_DO_NOT_COPY_VALUES, &th->MatIntegralSensip0));
    993:   }
    994:   /* backup sensitivity results for roll-backs */
    995:   PetscCall(MatDuplicate(ts->mat_sensip, MAT_DO_NOT_COPY_VALUES, &th->MatFwdSensip0));

    997:   PetscCall(VecDuplicate(ts->vec_sol, &th->VecDeltaFwdSensipCol));
    998:   PetscFunctionReturn(PETSC_SUCCESS);
    999: }

    1001: static PetscErrorCode TSForwardReset_Theta(TS ts)
    1002: {
    1003:   TS_Theta *th     = (TS_Theta *)ts->data;
    1004:   TS        quadts = ts->quadraturets;

    1006:   PetscFunctionBegin;
    1007:   if (quadts && quadts->mat_sensip) {
    1008:     PetscCall(MatDestroy(&th->MatIntegralSensipTemp));
    1009:     PetscCall(MatDestroy(&th->MatIntegralSensip0));
    1010:   }
    1011:   PetscCall(VecDestroy(&th->VecDeltaFwdSensipCol));
    1012:   PetscCall(MatDestroy(&th->MatDeltaFwdSensip));
    1013:   PetscCall(MatDestroy(&th->MatFwdSensip0));
    1014:   PetscFunctionReturn(PETSC_SUCCESS);
    1015: }

    1017: static PetscErrorCode TSSetUp_Theta(TS ts)
    1018: {
    1019:   TS_Theta *th     = (TS_Theta *)ts->data;
    1020:   TS        quadts = ts->quadraturets;
    1021:   PetscBool match;

    1023:   PetscFunctionBegin;
    1024:   if (!th->VecCostIntegral0 && quadts && ts->costintegralfwd) { /* back up cost integral */
    1025:     PetscCall(VecDuplicate(quadts->vec_sol, &th->VecCostIntegral0));
    1026:   }
    1027:   if (!th->X) PetscCall(VecDuplicate(ts->vec_sol, &th->X));
    1028:   if (!th->Xdot) PetscCall(VecDuplicate(ts->vec_sol, &th->Xdot));
    1029:   if (!th->X0) PetscCall(VecDuplicate(ts->vec_sol, &th->X0));
    1030:   if (th->endpoint) PetscCall(VecDuplicate(ts->vec_sol, &th->affine));

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

    1035:   PetscCall(TSGetDM(ts, &ts->dm));
    1036:   PetscCall(DMCoarsenHookAdd(ts->dm, DMCoarsenHook_TSTheta, DMRestrictHook_TSTheta, ts));
    1037:   PetscCall(DMSubDomainHookAdd(ts->dm, DMSubDomainHook_TSTheta, DMSubDomainRestrictHook_TSTheta, ts));

    1039:   PetscCall(TSGetAdapt(ts, &ts->adapt));
    1040:   PetscCall(TSAdaptCandidatesClear(ts->adapt));
    1041:   PetscCall(PetscObjectTypeCompare((PetscObject)ts->adapt, TSADAPTNONE, &match));
    1042:   if (!match) {
    1043:     if (!th->vec_sol_prev) PetscCall(VecDuplicate(ts->vec_sol, &th->vec_sol_prev));
    1044:     if (!th->vec_lte_work) PetscCall(VecDuplicate(ts->vec_sol, &th->vec_lte_work));
    1045:   }
    1046:   PetscCall(TSGetSNES(ts, &ts->snes));

    1048:   ts->stifflyaccurate = (!th->endpoint && th->Theta != 1.0) ? PETSC_FALSE : PETSC_TRUE;
    1049:   PetscFunctionReturn(PETSC_SUCCESS);
    1050: }

    1052: static PetscErrorCode TSAdjointSetUp_Theta(TS ts)
    1053: {
    1054:   TS_Theta *th = (TS_Theta *)ts->data;

    1056:   PetscFunctionBegin;
    1057:   PetscCall(VecDuplicateVecs(ts->vecs_sensi[0], ts->numcost, &th->VecsDeltaLam));
    1058:   PetscCall(VecDuplicateVecs(ts->vecs_sensi[0], ts->numcost, &th->VecsSensiTemp));
    1059:   if (ts->vecs_sensip) PetscCall(VecDuplicateVecs(ts->vecs_sensip[0], ts->numcost, &th->VecsDeltaMu));
    1060:   if (ts->vecs_sensi2) {
    1061:     PetscCall(VecDuplicateVecs(ts->vecs_sensi[0], ts->numcost, &th->VecsDeltaLam2));
    1062:     PetscCall(VecDuplicateVecs(ts->vecs_sensi2[0], ts->numcost, &th->VecsSensi2Temp));
    1063:     /* hack ts to make implicit TS solver work when users provide only explicit versions of callbacks (RHSFunction,RHSJacobian,RHSHessian etc.) */
    1064:     if (!ts->ihessianproduct_fuu) ts->vecs_fuu = ts->vecs_guu;
    1065:     if (!ts->ihessianproduct_fup) ts->vecs_fup = ts->vecs_gup;
    1066:   }
    1067:   if (ts->vecs_sensi2p) {
    1068:     PetscCall(VecDuplicateVecs(ts->vecs_sensi2p[0], ts->numcost, &th->VecsDeltaMu2));
    1069:     /* hack ts to make implicit TS solver work when users provide only explicit versions of callbacks (RHSFunction,RHSJacobian,RHSHessian etc.) */
    1070:     if (!ts->ihessianproduct_fpu) ts->vecs_fpu = ts->vecs_gpu;
    1071:     if (!ts->ihessianproduct_fpp) ts->vecs_fpp = ts->vecs_gpp;
    1072:   }
    1073:   PetscFunctionReturn(PETSC_SUCCESS);
    1074: }

    1076: static PetscErrorCode TSSetFromOptions_Theta(TS ts, PetscOptionItems PetscOptionsObject)
    1077: {
    1078:   TS_Theta *th = (TS_Theta *)ts->data;

    1080:   PetscFunctionBegin;
    1081:   PetscOptionsHeadBegin(PetscOptionsObject, "Theta ODE solver options");
    1082:   {
    1083:     PetscCall(PetscOptionsReal("-ts_theta_theta", "Location of stage (0<Theta<=1)", "TSThetaSetTheta", th->Theta, &th->Theta, NULL));
    1084:     PetscCall(PetscOptionsBool("-ts_theta_endpoint", "Use the endpoint instead of midpoint form of the Theta method", "TSThetaSetEndpoint", th->endpoint, &th->endpoint, NULL));
    1085:     PetscCall(PetscOptionsBool("-ts_theta_initial_guess_extrapolate", "Extrapolate stage initial guess from previous solution (sometimes unstable)", "TSThetaSetExtrapolate", th->extrapolate, &th->extrapolate, NULL));
    1086:   }
    1087:   PetscOptionsHeadEnd();
    1088:   PetscFunctionReturn(PETSC_SUCCESS);
    1089: }

    1091: static PetscErrorCode TSView_Theta(TS ts, PetscViewer viewer)
    1092: {
    1093:   TS_Theta *th = (TS_Theta *)ts->data;
    1094:   PetscBool isascii;

    1096:   PetscFunctionBegin;
    1097:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
    1098:   if (isascii) {
    1099:     PetscCall(PetscViewerASCIIPrintf(viewer, "  Theta=%g\n", (double)th->Theta));
    1100:     PetscCall(PetscViewerASCIIPrintf(viewer, "  Extrapolation=%s\n", th->extrapolate ? "yes" : "no"));
    1101:   }
    1102:   PetscFunctionReturn(PETSC_SUCCESS);
    1103: }

    1105: static PetscErrorCode TSThetaGetTheta_Theta(TS ts, PetscReal *theta)
    1106: {
    1107:   TS_Theta *th = (TS_Theta *)ts->data;

    1109:   PetscFunctionBegin;
    1110:   *theta = th->Theta;
    1111:   PetscFunctionReturn(PETSC_SUCCESS);
    1112: }

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

    1118:   PetscFunctionBegin;
    1119:   PetscCheck(theta > 0 && theta <= 1, PetscObjectComm((PetscObject)ts), PETSC_ERR_ARG_OUTOFRANGE, "Theta %g not in range (0,1]", (double)theta);
    1120:   th->Theta = theta;
    1121:   th->order = (th->Theta == 0.5) ? 2 : 1;
    1122:   PetscFunctionReturn(PETSC_SUCCESS);
    1123: }

    1125: static PetscErrorCode TSThetaGetEndpoint_Theta(TS ts, PetscBool *endpoint)
    1126: {
    1127:   TS_Theta *th = (TS_Theta *)ts->data;

    1129:   PetscFunctionBegin;
    1130:   *endpoint = th->endpoint;
    1131:   PetscFunctionReturn(PETSC_SUCCESS);
    1132: }

    1134: static PetscErrorCode TSThetaSetEndpoint_Theta(TS ts, PetscBool flg)
    1135: {
    1136:   TS_Theta *th = (TS_Theta *)ts->data;

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

    1143: #if defined(PETSC_HAVE_COMPLEX)
    1144: static PetscErrorCode TSComputeLinearStability_Theta(TS ts, PetscReal xr, PetscReal xi, PetscReal *yr, PetscReal *yi)
    1145: {
    1146:   PetscComplex z  = xr + xi * PETSC_i, f;
    1147:   TS_Theta    *th = (TS_Theta *)ts->data;

    1149:   PetscFunctionBegin;
    1150:   f   = (1.0 + (1.0 - th->Theta) * z) / (1.0 - th->Theta * z);
    1151:   *yr = PetscRealPartComplex(f);
    1152:   *yi = PetscImaginaryPartComplex(f);
    1153:   PetscFunctionReturn(PETSC_SUCCESS);
    1154: }
    1155: #endif

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

    1161:   PetscFunctionBegin;
    1162:   if (ns) {
    1163:     if (!th->endpoint && th->Theta != 1.0) *ns = 1; /* midpoint form */
    1164:     else *ns = 2;                                   /* endpoint form */
    1165:   }
    1166:   if (Y) {
    1167:     if (!th->endpoint && th->Theta != 1.0) {
    1168:       th->Stages[0] = th->X;
    1169:     } else {
    1170:       th->Stages[0] = th->X0;
    1171:       th->Stages[1] = ts->vec_sol; /* stiffly accurate */
    1172:     }
    1173:     *Y = th->Stages;
    1174:   }
    1175:   PetscFunctionReturn(PETSC_SUCCESS);
    1176: }

    1178: /*MC
    1179:       TSTHETA - DAE solver using the implicit Theta method

    1181:    Level: beginner

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

    1188:    Notes:
    1189: .vb
    1190:   -ts_type theta -ts_theta_theta 1.0 corresponds to backward Euler (TSBEULER)
    1191:   -ts_type theta -ts_theta_theta 0.5 corresponds to the implicit midpoint rule
    1192:   -ts_type theta -ts_theta_theta 0.5 -ts_theta_endpoint corresponds to Crank-Nicholson (TSCN)
    1193: .ve

    1195:    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.

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

    1199: .vb
    1200:   Theta | Theta
    1201:   -------------
    1202:         |  1
    1203: .ve

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

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

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

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

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

    1224: .seealso: [](ch_ts), `TSCreate()`, `TS`, `TSSetType()`, `TSCN`, `TSBEULER`, `TSThetaSetTheta()`, `TSThetaSetEndpoint()`
    1225: M*/
    1226: PETSC_EXTERN PetscErrorCode TSCreate_Theta(TS ts)
    1227: {
    1228:   TS_Theta *th;

    1230:   PetscFunctionBegin;
    1231:   ts->ops->reset          = TSReset_Theta;
    1232:   ts->ops->adjointreset   = TSAdjointReset_Theta;
    1233:   ts->ops->destroy        = TSDestroy_Theta;
    1234:   ts->ops->view           = TSView_Theta;
    1235:   ts->ops->setup          = TSSetUp_Theta;
    1236:   ts->ops->adjointsetup   = TSAdjointSetUp_Theta;
    1237:   ts->ops->adjointreset   = TSAdjointReset_Theta;
    1238:   ts->ops->step           = TSStep_Theta;
    1239:   ts->ops->interpolate    = TSInterpolate_Theta;
    1240:   ts->ops->evaluatewlte   = TSEvaluateWLTE_Theta;
    1241:   ts->ops->rollback       = TSRollBack_Theta;
    1242:   ts->ops->resizeregister = TSResizeRegister_Theta;
    1243:   ts->ops->setfromoptions = TSSetFromOptions_Theta;
    1244:   ts->ops->snesfunction   = SNESTSFormFunction_Theta;
    1245:   ts->ops->snesjacobian   = SNESTSFormJacobian_Theta;
    1246: #if defined(PETSC_HAVE_COMPLEX)
    1247:   ts->ops->linearstability = TSComputeLinearStability_Theta;
    1248: #endif
    1249:   ts->ops->getstages       = TSGetStages_Theta;
    1250:   ts->ops->adjointstep     = TSAdjointStep_Theta;
    1251:   ts->ops->adjointintegral = TSAdjointCostIntegral_Theta;
    1252:   ts->ops->forwardintegral = TSForwardCostIntegral_Theta;
    1253:   ts->default_adapt_type   = TSADAPTNONE;

    1255:   ts->ops->forwardsetup     = TSForwardSetUp_Theta;
    1256:   ts->ops->forwardreset     = TSForwardReset_Theta;
    1257:   ts->ops->forwardstep      = TSForwardStep_Theta;
    1258:   ts->ops->forwardgetstages = TSForwardGetStages_Theta;

    1260:   ts->usessnes = PETSC_TRUE;

    1262:   PetscCall(PetscNew(&th));
    1263:   ts->data = (void *)th;

    1265:   th->VecsDeltaLam   = NULL;
    1266:   th->VecsDeltaMu    = NULL;
    1267:   th->VecsSensiTemp  = NULL;
    1268:   th->VecsSensi2Temp = NULL;

    1270:   th->extrapolate = PETSC_FALSE;
    1271:   th->Theta       = 0.5;
    1272:   th->order       = 2;
    1273:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaGetTheta_C", TSThetaGetTheta_Theta));
    1274:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaSetTheta_C", TSThetaSetTheta_Theta));
    1275:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaGetEndpoint_C", TSThetaGetEndpoint_Theta));
    1276:   PetscCall(PetscObjectComposeFunction((PetscObject)ts, "TSThetaSetEndpoint_C", TSThetaSetEndpoint_Theta));
    1277:   PetscFunctionReturn(PETSC_SUCCESS);
    1278: }

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

    1283:   Not Collective

    1285:   Input Parameter:
    1286: . ts - timestepping context

    1288:   Output Parameter:
    1289: . theta - stage abscissa

    1291:   Level: advanced

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

    1296: .seealso: [](ch_ts), `TSThetaSetTheta()`, `TSTHETA`
    1297: @*/
    1298: PetscErrorCode TSThetaGetTheta(TS ts, PetscReal *theta)
    1299: {
    1300:   PetscFunctionBegin;
    1302:   PetscAssertPointer(theta, 2);
    1303:   PetscUseMethod(ts, "TSThetaGetTheta_C", (TS, PetscReal *), (ts, theta));
    1304:   PetscFunctionReturn(PETSC_SUCCESS);
    1305: }

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

    1310:   Not Collective

    1312:   Input Parameters:
    1313: + ts    - timestepping context
    1314: - theta - stage abscissa

    1316:   Options Database Key:
    1317: . -ts_theta_theta theta - set theta

    1319:   Level: intermediate

    1321: .seealso: [](ch_ts), `TSThetaGetTheta()`, `TSTHETA`, `TSCN`
    1322: @*/
    1323: PetscErrorCode TSThetaSetTheta(TS ts, PetscReal theta)
    1324: {
    1325:   PetscFunctionBegin;
    1327:   PetscTryMethod(ts, "TSThetaSetTheta_C", (TS, PetscReal), (ts, theta));
    1328:   PetscFunctionReturn(PETSC_SUCCESS);
    1329: }

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

    1334:   Not Collective

    1336:   Input Parameter:
    1337: . ts - timestepping context

    1339:   Output Parameter:
    1340: . endpoint - `PETSC_TRUE` when using the endpoint variant

    1342:   Level: advanced

    1344: .seealso: [](ch_ts), `TSThetaSetEndpoint()`, `TSTHETA`, `TSCN`
    1345: @*/
    1346: PetscErrorCode TSThetaGetEndpoint(TS ts, PetscBool *endpoint)
    1347: {
    1348:   PetscFunctionBegin;
    1350:   PetscAssertPointer(endpoint, 2);
    1351:   PetscUseMethod(ts, "TSThetaGetEndpoint_C", (TS, PetscBool *), (ts, endpoint));
    1352:   PetscFunctionReturn(PETSC_SUCCESS);
    1353: }

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

    1358:   Not Collective

    1360:   Input Parameters:
    1361: + ts  - timestepping context
    1362: - flg - `PETSC_TRUE` to use the endpoint variant

    1364:   Options Database Key:
    1365: . -ts_theta_endpoint flg - use the endpoint variant

    1367:   Level: intermediate

    1369: .seealso: [](ch_ts), `TSTHETA`, `TSCN`
    1370: @*/
    1371: PetscErrorCode TSThetaSetEndpoint(TS ts, PetscBool flg)
    1372: {
    1373:   PetscFunctionBegin;
    1375:   PetscTryMethod(ts, "TSThetaSetEndpoint_C", (TS, PetscBool), (ts, flg));
    1376:   PetscFunctionReturn(PETSC_SUCCESS);
    1377: }

    1379: /*
    1380:  * TSBEULER and TSCN are straightforward specializations of TSTHETA.
    1381:  * The creation functions for these specializations are below.
    1382:  */

    1384: static PetscErrorCode TSSetUp_BEuler(TS ts)
    1385: {
    1386:   TS_Theta *th = (TS_Theta *)ts->data;

    1388:   PetscFunctionBegin;
    1389:   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");
    1390:   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");
    1391:   PetscCall(TSSetUp_Theta(ts));
    1392:   PetscFunctionReturn(PETSC_SUCCESS);
    1393: }

    1395: static PetscErrorCode TSView_BEuler(TS ts, PetscViewer viewer)
    1396: {
    1397:   PetscFunctionBegin;
    1398:   PetscFunctionReturn(PETSC_SUCCESS);
    1399: }

    1401: /*MC
    1402:       TSBEULER - ODE solver using the implicit backward Euler method

    1404:   Level: beginner

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

    1409: .seealso: [](ch_ts), `TSCreate()`, `TS`, `TSSetType()`, `TSEULER`, `TSCN`, `TSTHETA`
    1410: M*/
    1411: PETSC_EXTERN PetscErrorCode TSCreate_BEuler(TS ts)
    1412: {
    1413:   PetscFunctionBegin;
    1414:   PetscCall(TSCreate_Theta(ts));
    1415:   PetscCall(TSThetaSetTheta(ts, 1.0));
    1416:   PetscCall(TSThetaSetEndpoint(ts, PETSC_FALSE));
    1417:   ts->ops->setup = TSSetUp_BEuler;
    1418:   ts->ops->view  = TSView_BEuler;
    1419:   PetscFunctionReturn(PETSC_SUCCESS);
    1420: }

    1422: static PetscErrorCode TSSetUp_CN(TS ts)
    1423: {
    1424:   TS_Theta *th = (TS_Theta *)ts->data;

    1426:   PetscFunctionBegin;
    1427:   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");
    1428:   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");
    1429:   PetscCall(TSSetUp_Theta(ts));
    1430:   PetscFunctionReturn(PETSC_SUCCESS);
    1431: }

    1433: static PetscErrorCode TSView_CN(TS ts, PetscViewer viewer)
    1434: {
    1435:   PetscFunctionBegin;
    1436:   PetscFunctionReturn(PETSC_SUCCESS);
    1437: }

    1439: /*MC
    1440:   TSCN - ODE solver using the implicit Crank-Nicolson method.

    1442:   Level: beginner

    1444:   Notes:
    1445:   `TSCN` is equivalent to `TSTHETA` with Theta=0.5 and the "endpoint" option set. I.e.
    1446: .vb
    1447:   -ts_type theta
    1448:   -ts_theta_theta 0.5
    1449:   -ts_theta_endpoint
    1450: .ve

    1452: .seealso: [](ch_ts), `TSCreate()`, `TS`, `TSSetType()`, `TSBEULER`, `TSTHETA`, `TSType`
    1453: M*/
    1454: PETSC_EXTERN PetscErrorCode TSCreate_CN(TS ts)
    1455: {
    1456:   PetscFunctionBegin;
    1457:   PetscCall(TSCreate_Theta(ts));
    1458:   PetscCall(TSThetaSetTheta(ts, 0.5));
    1459:   PetscCall(TSThetaSetEndpoint(ts, PETSC_TRUE));
    1460:   ts->ops->setup = TSSetUp_CN;
    1461:   ts->ops->view  = TSView_CN;
    1462:   PetscFunctionReturn(PETSC_SUCCESS);
    1463: }