implicit/theta/theta.c

/*
  Code for timestepping with implicit Theta method
*/
#include <petsc/private/tsimpl.h>                /*I   "petscts.h"   I*/
#include <petscsnes.h>
#include <petscdm.h>
#include <petscmat.h>

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

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

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

  PetscFunctionBegin;
  if (X0) {
    if (dm && dm != ts->dm) {
      ierr = DMGetNamedGlobalVector(dm,"TSTheta_X0",X0);CHKERRQ(ierr);
    } else *X0 = ts->vec_sol;
  }
  if (Xdot) {
    if (dm && dm != ts->dm) {
      ierr = DMGetNamedGlobalVector(dm,"TSTheta_Xdot",Xdot);CHKERRQ(ierr);
    } else *Xdot = th->Xdot;
  }
  PetscFunctionReturn(0);
}

static PetscErrorCode TSThetaRestoreX0AndXdot(TS ts,DM dm,Vec *X0,Vec *Xdot)
{
  PetscErrorCode ierr;

  PetscFunctionBegin;
  if (X0) {
    if (dm && dm != ts->dm) {
      ierr = DMRestoreNamedGlobalVector(dm,"TSTheta_X0",X0);CHKERRQ(ierr);
    }
  }
  if (Xdot) {
    if (dm && dm != ts->dm) {
      ierr = DMRestoreNamedGlobalVector(dm,"TSTheta_Xdot",Xdot);CHKERRQ(ierr);
    }
  }
  PetscFunctionReturn(0);
}

static PetscErrorCode DMCoarsenHook_TSTheta(DM fine,DM coarse,void *ctx)
{
  PetscFunctionBegin;
  PetscFunctionReturn(0);
}

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

  PetscFunctionBegin;
  ierr = TSThetaGetX0AndXdot(ts,fine,&X0,&Xdot);CHKERRQ(ierr);
  ierr = TSThetaGetX0AndXdot(ts,coarse,&X0_c,&Xdot_c);CHKERRQ(ierr);
  ierr = MatRestrict(restrct,X0,X0_c);CHKERRQ(ierr);
  ierr = MatRestrict(restrct,Xdot,Xdot_c);CHKERRQ(ierr);
  ierr = VecPointwiseMult(X0_c,rscale,X0_c);CHKERRQ(ierr);
  ierr = VecPointwiseMult(Xdot_c,rscale,Xdot_c);CHKERRQ(ierr);
  ierr = TSThetaRestoreX0AndXdot(ts,fine,&X0,&Xdot);CHKERRQ(ierr);
  ierr = TSThetaRestoreX0AndXdot(ts,coarse,&X0_c,&Xdot_c);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

static PetscErrorCode DMSubDomainHook_TSTheta(DM dm,DM subdm,void *ctx)
{
  PetscFunctionBegin;
  PetscFunctionReturn(0);
}

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

  PetscFunctionBegin;
  ierr = TSThetaGetX0AndXdot(ts,dm,&X0,&Xdot);CHKERRQ(ierr);
  ierr = TSThetaGetX0AndXdot(ts,subdm,&X0_sub,&Xdot_sub);CHKERRQ(ierr);

  ierr = VecScatterBegin(gscat,X0,X0_sub,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr);
  ierr = VecScatterEnd(gscat,X0,X0_sub,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr);

  ierr = VecScatterBegin(gscat,Xdot,Xdot_sub,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr);
  ierr = VecScatterEnd(gscat,Xdot,Xdot_sub,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr);

  ierr = TSThetaRestoreX0AndXdot(ts,dm,&X0,&Xdot);CHKERRQ(ierr);
  ierr = TSThetaRestoreX0AndXdot(ts,subdm,&X0_sub,&Xdot_sub);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

static PetscErrorCode TSThetaEvaluateCostIntegral(TS ts)
{
  TS_Theta       *th = (TS_Theta*)ts->data;
  TS             quadts = ts->quadraturets;
  PetscErrorCode ierr;

  PetscFunctionBegin;
  if (th->endpoint) {
    /* Evolve ts->vec_costintegral to compute integrals */
    if (th->Theta!=1.0) {
      ierr = TSComputeRHSFunction(quadts,th->ptime0,th->X0,ts->vec_costintegrand);CHKERRQ(ierr);
      ierr = VecAXPY(quadts->vec_sol,th->time_step0*(1.0-th->Theta),ts->vec_costintegrand);CHKERRQ(ierr);
    }
    ierr = TSComputeRHSFunction(quadts,ts->ptime,ts->vec_sol,ts->vec_costintegrand);CHKERRQ(ierr);
    ierr = VecAXPY(quadts->vec_sol,th->time_step0*th->Theta,ts->vec_costintegrand);CHKERRQ(ierr);
  } else {
    ierr = TSComputeRHSFunction(quadts,th->stage_time,th->X,ts->vec_costintegrand);CHKERRQ(ierr);
    ierr = VecAXPY(quadts->vec_sol,th->time_step0,ts->vec_costintegrand);CHKERRQ(ierr);
  }
  PetscFunctionReturn(0);
}

static PetscErrorCode TSForwardCostIntegral_Theta(TS ts)
{
  TS_Theta       *th = (TS_Theta*)ts->data;
  TS             quadts = ts->quadraturets;
  PetscErrorCode ierr;

  PetscFunctionBegin;
  /* backup cost integral */
  ierr = VecCopy(quadts->vec_sol,th->VecCostIntegral0);CHKERRQ(ierr);
  ierr = TSThetaEvaluateCostIntegral(ts);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

static PetscErrorCode TSAdjointCostIntegral_Theta(TS ts)
{
  TS_Theta       *th = (TS_Theta*)ts->data;
  PetscErrorCode ierr;

  PetscFunctionBegin;
  /* Like TSForwardCostIntegral(), the adjoint cost integral evaluation relies on ptime0 and time_step0. */
  th->ptime0     = ts->ptime + ts->time_step;
  th->time_step0 = -ts->time_step;
  ierr = TSThetaEvaluateCostIntegral(ts);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

static PetscErrorCode TSTheta_SNESSolve(TS ts,Vec b,Vec x)
{
  PetscInt       nits,lits;
  PetscErrorCode ierr;

  PetscFunctionBegin;
  ierr = SNESSolve(ts->snes,b,x);CHKERRQ(ierr);
  ierr = SNESGetIterationNumber(ts->snes,&nits);CHKERRQ(ierr);
  ierr = SNESGetLinearSolveIterations(ts->snes,&lits);CHKERRQ(ierr);
  ts->snes_its += nits; ts->ksp_its += lits;
  PetscFunctionReturn(0);
}

static PetscErrorCode TSStep_Theta(TS ts)
{
  TS_Theta       *th = (TS_Theta*)ts->data;
  PetscInt       rejections = 0;
  PetscBool      stageok,accept = PETSC_TRUE;
  PetscReal      next_time_step = ts->time_step;
  PetscErrorCode ierr;

  PetscFunctionBegin;
  if (!ts->steprollback) {
    if (th->vec_sol_prev) { ierr = VecCopy(th->X0,th->vec_sol_prev);CHKERRQ(ierr); }
    ierr = VecCopy(ts->vec_sol,th->X0);CHKERRQ(ierr);
  }

  th->status     = TS_STEP_INCOMPLETE;
  while (!ts->reason && th->status != TS_STEP_COMPLETE) {
    th->shift      = 1/(th->Theta*ts->time_step);
    th->stage_time = ts->ptime + (th->endpoint ? (PetscReal)1 : th->Theta)*ts->time_step;
    ierr = VecCopy(th->X0,th->X);CHKERRQ(ierr);
    if (th->extrapolate && !ts->steprestart) {
      ierr = VecAXPY(th->X,1/th->shift,th->Xdot);CHKERRQ(ierr);
    }
    if (th->endpoint) { /* This formulation assumes linear time-independent mass matrix */
      if (!th->affine) {ierr = VecDuplicate(ts->vec_sol,&th->affine);CHKERRQ(ierr);}
      ierr = VecZeroEntries(th->Xdot);CHKERRQ(ierr);
      ierr = TSComputeIFunction(ts,ts->ptime,th->X0,th->Xdot,th->affine,PETSC_FALSE);CHKERRQ(ierr);
      ierr = VecScale(th->affine,(th->Theta-1)/th->Theta);CHKERRQ(ierr);
    } else if (th->affine) { /* Just in case th->endpoint is changed between calls to TSStep_Theta() */
      ierr = VecZeroEntries(th->affine);CHKERRQ(ierr);
    }
    ierr = TSPreStage(ts,th->stage_time);CHKERRQ(ierr);
    ierr = TSTheta_SNESSolve(ts,th->affine,th->X);CHKERRQ(ierr);
    ierr = TSPostStage(ts,th->stage_time,0,&th->X);CHKERRQ(ierr);
    ierr = TSAdaptCheckStage(ts->adapt,ts,th->stage_time,th->X,&stageok);CHKERRQ(ierr);
    if (!stageok) goto reject_step;

    th->status = TS_STEP_PENDING;
    if (th->endpoint) {
      ierr = VecCopy(th->X,ts->vec_sol);CHKERRQ(ierr);
    } else {
      ierr = VecAXPBYPCZ(th->Xdot,-th->shift,th->shift,0,th->X0,th->X);CHKERRQ(ierr);
      ierr = VecAXPY(ts->vec_sol,ts->time_step,th->Xdot);CHKERRQ(ierr);
    }
    ierr = TSAdaptChoose(ts->adapt,ts,ts->time_step,NULL,&next_time_step,&accept);CHKERRQ(ierr);
    th->status = accept ? TS_STEP_COMPLETE : TS_STEP_INCOMPLETE;
    if (!accept) {
      ierr = VecCopy(th->X0,ts->vec_sol);CHKERRQ(ierr);
      ts->time_step = next_time_step;
      goto reject_step;
    }

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

  reject_step:
    ts->reject++; accept = PETSC_FALSE;
    if (!ts->reason && ++rejections > ts->max_reject && ts->max_reject >= 0) {
      ts->reason = TS_DIVERGED_STEP_REJECTED;
      ierr = PetscInfo2(ts,"Step=%D, step rejections %D greater than current TS allowed, stopping solve\n",ts->steps,rejections);CHKERRQ(ierr);
    }
  }
  PetscFunctionReturn(0);
}

static PetscErrorCode TSAdjointStepBEuler_Private(TS ts)
{
  TS_Theta       *th = (TS_Theta*)ts->data;
  TS             quadts = ts->quadraturets;
  Vec            *VecsDeltaLam = th->VecsDeltaLam,*VecsDeltaMu = th->VecsDeltaMu,*VecsSensiTemp = th->VecsSensiTemp;
  Vec            *VecsDeltaLam2 = th->VecsDeltaLam2,*VecsDeltaMu2 = th->VecsDeltaMu2,*VecsSensi2Temp = th->VecsSensi2Temp;
  PetscInt       nadj;
  Mat            J,Jpre,quadJ = NULL,quadJp = NULL;
  KSP            ksp;
  PetscScalar    *xarr;
  TSEquationType eqtype;
  PetscBool      isexplicitode = PETSC_FALSE;
  PetscReal      adjoint_time_step;
  PetscErrorCode ierr;

  PetscFunctionBegin;
  ierr = TSGetEquationType(ts,&eqtype);CHKERRQ(ierr);
  if (eqtype == TS_EQ_ODE_EXPLICIT) {
    isexplicitode  = PETSC_TRUE;
    VecsDeltaLam  = ts->vecs_sensi;
    VecsDeltaLam2 = ts->vecs_sensi2;
  }
  th->status = TS_STEP_INCOMPLETE;
  ierr = SNESGetKSP(ts->snes,&ksp);CHKERRQ(ierr);
  ierr = TSGetIJacobian(ts,&J,&Jpre,NULL,NULL);CHKERRQ(ierr);
  if (quadts) {
    ierr = TSGetRHSJacobian(quadts,&quadJ,NULL,NULL,NULL);CHKERRQ(ierr);
    ierr = TSGetRHSJacobianP(quadts,&quadJp,NULL,NULL);CHKERRQ(ierr);
  }

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

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

  for (nadj=0; nadj<ts->numcost; nadj++) {
    ierr = VecCopy(ts->vecs_sensi[nadj],VecsSensiTemp[nadj]);CHKERRQ(ierr);
    ierr = VecScale(VecsSensiTemp[nadj],1./adjoint_time_step);CHKERRQ(ierr); /* lambda_{n+1}/h */
    if (quadJ) {
      ierr = MatDenseGetColumn(quadJ,nadj,&xarr);CHKERRQ(ierr);
      ierr = VecPlaceArray(ts->vec_drdu_col,xarr);CHKERRQ(ierr);
      ierr = VecAXPY(VecsSensiTemp[nadj],1.,ts->vec_drdu_col);CHKERRQ(ierr);
      ierr = VecResetArray(ts->vec_drdu_col);CHKERRQ(ierr);
      ierr = MatDenseRestoreColumn(quadJ,&xarr);CHKERRQ(ierr);
    }
  }

  /* Build LHS for first-order adjoint */
  th->shift = 1./adjoint_time_step;
  ierr = TSComputeSNESJacobian(ts,ts->vec_sol,J,Jpre);CHKERRQ(ierr);
  ierr = KSPSetOperators(ksp,J,Jpre);CHKERRQ(ierr);

  /* Solve stage equation LHS*lambda_s = RHS for first-order adjoint */
  for (nadj=0; nadj<ts->numcost; nadj++) {
    KSPConvergedReason kspreason;
    ierr = KSPSolveTranspose(ksp,VecsSensiTemp[nadj],VecsDeltaLam[nadj]);CHKERRQ(ierr);
    ierr = KSPGetConvergedReason(ksp,&kspreason);CHKERRQ(ierr);
    if (kspreason < 0) {
      ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
      ierr = PetscInfo2(ts,"Step=%D, %Dth cost function, transposed linear solve fails, stopping 1st-order adjoint solve\n",ts->steps,nadj);CHKERRQ(ierr);
    }
  }

  if (ts->vecs_sensi2) { /* U_{n+1} */
    /* Get w1 at t_{n+1} from TLM matrix */
    ierr = MatDenseGetColumn(ts->mat_sensip,0,&xarr);CHKERRQ(ierr);
    ierr = VecPlaceArray(ts->vec_sensip_col,xarr);CHKERRQ(ierr);
    /* lambda_s^T F_UU w_1 */
    ierr = TSComputeIHessianProductFunctionUU(ts,th->stage_time,ts->vec_sol,VecsDeltaLam,ts->vec_sensip_col,ts->vecs_fuu);CHKERRQ(ierr);
    /* lambda_s^T F_UP w_2 */
    ierr = TSComputeIHessianProductFunctionUP(ts,th->stage_time,ts->vec_sol,VecsDeltaLam,ts->vec_dir,ts->vecs_fup);CHKERRQ(ierr);
    for (nadj=0; nadj<ts->numcost; nadj++) { /* compute the residual */
      ierr = VecCopy(ts->vecs_sensi2[nadj],VecsSensi2Temp[nadj]);CHKERRQ(ierr);
      ierr = VecScale(VecsSensi2Temp[nadj],1./adjoint_time_step);CHKERRQ(ierr);
      ierr = VecAXPY(VecsSensi2Temp[nadj],-1.,ts->vecs_fuu[nadj]);CHKERRQ(ierr);
      if (ts->vecs_fup) {
        ierr = VecAXPY(VecsSensi2Temp[nadj],-1.,ts->vecs_fup[nadj]);CHKERRQ(ierr);
      }
    }
    /* Solve stage equation LHS X = RHS for second-order adjoint */
    for (nadj=0; nadj<ts->numcost; nadj++) {
      KSPConvergedReason kspreason;
      ierr = KSPSolveTranspose(ksp,VecsSensi2Temp[nadj],VecsDeltaLam2[nadj]);CHKERRQ(ierr);
      ierr = KSPGetConvergedReason(ksp,&kspreason);CHKERRQ(ierr);
      if (kspreason < 0) {
        ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
        ierr = PetscInfo2(ts,"Step=%D, %Dth cost function, transposed linear solve fails, stopping 2nd-order adjoint solve\n",ts->steps,nadj);CHKERRQ(ierr);
      }
    }
  }

  /* Update sensitivities, and evaluate integrals if there is any */
  if (!isexplicitode) {
    th->shift = 0.0;
    ierr = TSComputeSNESJacobian(ts,ts->vec_sol,J,Jpre);CHKERRQ(ierr);
    ierr = KSPSetOperators(ksp,J,Jpre);CHKERRQ(ierr);
    ierr = MatScale(J,-1.);CHKERRQ(ierr);
    for (nadj=0; nadj<ts->numcost; nadj++) {
      /* Add f_U \lambda_s to the original RHS */
      ierr = MatMultTransposeAdd(J,VecsDeltaLam[nadj],VecsSensiTemp[nadj],VecsSensiTemp[nadj]);CHKERRQ(ierr);
      ierr = VecScale(VecsSensiTemp[nadj],adjoint_time_step);CHKERRQ(ierr);
      ierr = VecCopy(VecsSensiTemp[nadj],ts->vecs_sensi[nadj]);CHKERRQ(ierr);
      if (ts->vecs_sensi2) {
        ierr = MatMultTransposeAdd(J,VecsDeltaLam2[nadj],VecsSensi2Temp[nadj],VecsSensi2Temp[nadj]);CHKERRQ(ierr);
        ierr = VecScale(VecsSensi2Temp[nadj],adjoint_time_step);CHKERRQ(ierr);
        ierr = VecCopy(VecsSensi2Temp[nadj],ts->vecs_sensi2[nadj]);CHKERRQ(ierr);
      }
    }
  }
  if (ts->vecs_sensip) {
    ierr = VecAXPBYPCZ(th->Xdot,-1./adjoint_time_step,1.0/adjoint_time_step,0,th->X0,ts->vec_sol);CHKERRQ(ierr);
    ierr = TSComputeIJacobianP(ts,th->stage_time,ts->vec_sol,th->Xdot,1./adjoint_time_step,ts->Jacp,PETSC_FALSE);CHKERRQ(ierr); /* get -f_p */
    if (quadts) {
      ierr = TSComputeRHSJacobianP(quadts,th->stage_time,ts->vec_sol,quadJp);CHKERRQ(ierr);
    }
    if (ts->vecs_sensi2p) {
      /* lambda_s^T F_PU w_1 */
      ierr = TSComputeIHessianProductFunctionPU(ts,th->stage_time,ts->vec_sol,VecsDeltaLam,ts->vec_sensip_col,ts->vecs_fpu);CHKERRQ(ierr);
      /* lambda_s^T F_PP w_2 */
      ierr = TSComputeIHessianProductFunctionPP(ts,th->stage_time,ts->vec_sol,VecsDeltaLam,ts->vec_dir,ts->vecs_fpp);CHKERRQ(ierr);
    }

    for (nadj=0; nadj<ts->numcost; nadj++) {
      ierr = MatMultTranspose(ts->Jacp,VecsDeltaLam[nadj],VecsDeltaMu[nadj]);CHKERRQ(ierr);
      ierr = VecAXPY(ts->vecs_sensip[nadj],-adjoint_time_step,VecsDeltaMu[nadj]);CHKERRQ(ierr);
      if (quadJp) {
        ierr = MatDenseGetColumn(quadJp,nadj,&xarr);CHKERRQ(ierr);
        ierr = VecPlaceArray(ts->vec_drdp_col,xarr);CHKERRQ(ierr);
        ierr = VecAXPY(ts->vecs_sensip[nadj],adjoint_time_step,ts->vec_drdp_col);CHKERRQ(ierr);
        ierr = VecResetArray(ts->vec_drdp_col);CHKERRQ(ierr);
        ierr = MatDenseRestoreColumn(quadJp,&xarr);CHKERRQ(ierr);
      }
      if (ts->vecs_sensi2p) {
        ierr = MatMultTranspose(ts->Jacp,VecsDeltaLam2[nadj],VecsDeltaMu2[nadj]);CHKERRQ(ierr);
        ierr = VecAXPY(ts->vecs_sensi2p[nadj],-adjoint_time_step,VecsDeltaMu2[nadj]);CHKERRQ(ierr);
        if (ts->vecs_fpu) {
          ierr = VecAXPY(ts->vecs_sensi2p[nadj],-adjoint_time_step,ts->vecs_fpu[nadj]);CHKERRQ(ierr);
        }
        if (ts->vecs_fpp) {
          ierr = VecAXPY(ts->vecs_sensi2p[nadj],-adjoint_time_step,ts->vecs_fpp[nadj]);CHKERRQ(ierr);
        }
      }
    }
  }

  if (ts->vecs_sensi2) {
    ierr = VecResetArray(ts->vec_sensip_col);CHKERRQ(ierr);
    ierr = MatDenseRestoreColumn(ts->mat_sensip,&xarr);CHKERRQ(ierr);
  }
  th->status = TS_STEP_COMPLETE;
  PetscFunctionReturn(0);
}

static PetscErrorCode TSAdjointStep_Theta(TS ts)
{
  TS_Theta       *th = (TS_Theta*)ts->data;
  TS             quadts = ts->quadraturets;
  Vec            *VecsDeltaLam = th->VecsDeltaLam,*VecsDeltaMu = th->VecsDeltaMu,*VecsSensiTemp = th->VecsSensiTemp;
  Vec            *VecsDeltaLam2 = th->VecsDeltaLam2,*VecsDeltaMu2 = th->VecsDeltaMu2,*VecsSensi2Temp = th->VecsSensi2Temp;
  PetscInt       nadj;
  Mat            J,Jpre,quadJ = NULL,quadJp = NULL;
  KSP            ksp;
  PetscScalar    *xarr;
  PetscReal      adjoint_time_step;
  PetscReal      adjoint_ptime; /* end time of the adjoint time step (ts->ptime is the start time, ususally ts->ptime is larger than adjoint_ptime) */
  PetscErrorCode ierr;

  PetscFunctionBegin;
  if (th->Theta == 1.) {
    ierr = TSAdjointStepBEuler_Private(ts);CHKERRQ(ierr);
    PetscFunctionReturn(0);
  }
  th->status = TS_STEP_INCOMPLETE;
  ierr = SNESGetKSP(ts->snes,&ksp);CHKERRQ(ierr);
  ierr = TSGetIJacobian(ts,&J,&Jpre,NULL,NULL);CHKERRQ(ierr);
  if (quadts) {
    ierr = TSGetRHSJacobian(quadts,&quadJ,NULL,NULL,NULL);CHKERRQ(ierr);
    ierr = TSGetRHSJacobianP(quadts,&quadJp,NULL,NULL);CHKERRQ(ierr);
  }
  /* If endpoint=1, th->ptime and th->X0 will be used; if endpoint=0, th->stage_time and th->X will be used. */
  th->stage_time    = th->endpoint ? ts->ptime : (ts->ptime+(1.-th->Theta)*ts->time_step);
  adjoint_ptime     = ts->ptime + ts->time_step;
  adjoint_time_step = -ts->time_step;  /* always positive since time_step is negative */

  if (!th->endpoint) {
    /* recover th->X0 using vec_sol and the stage value th->X */
    ierr = VecAXPBYPCZ(th->X0,1.0/(1.0-th->Theta),th->Theta/(th->Theta-1.0),0,th->X,ts->vec_sol);CHKERRQ(ierr);
  }

  /* Build RHS for first-order adjoint */
  /* Cost function has an integral term */
  if (quadts) {
    if (th->endpoint) {
      ierr = TSComputeRHSJacobian(quadts,th->stage_time,ts->vec_sol,quadJ,NULL);CHKERRQ(ierr);
    } else {
      ierr = TSComputeRHSJacobian(quadts,th->stage_time,th->X,quadJ,NULL);CHKERRQ(ierr);
    }
  }

  for (nadj=0; nadj<ts->numcost; nadj++) {
    ierr = VecCopy(ts->vecs_sensi[nadj],VecsSensiTemp[nadj]);CHKERRQ(ierr);
    ierr = VecScale(VecsSensiTemp[nadj],1./(th->Theta*adjoint_time_step));CHKERRQ(ierr);
    if (quadJ) {
      ierr = MatDenseGetColumn(quadJ,nadj,&xarr);CHKERRQ(ierr);
      ierr = VecPlaceArray(ts->vec_drdu_col,xarr);CHKERRQ(ierr);
      ierr = VecAXPY(VecsSensiTemp[nadj],1.,ts->vec_drdu_col);CHKERRQ(ierr);
      ierr = VecResetArray(ts->vec_drdu_col);CHKERRQ(ierr);
      ierr = MatDenseRestoreColumn(quadJ,&xarr);CHKERRQ(ierr);
    }
  }

  /* Build LHS for first-order adjoint */
  th->shift = 1./(th->Theta*adjoint_time_step);
  if (th->endpoint) {
    ierr = TSComputeSNESJacobian(ts,ts->vec_sol,J,Jpre);CHKERRQ(ierr);
  } else {
    ierr = TSComputeSNESJacobian(ts,th->X,J,Jpre);CHKERRQ(ierr);
  }
  ierr = KSPSetOperators(ksp,J,Jpre);CHKERRQ(ierr);

  /* Solve stage equation LHS*lambda_s = RHS for first-order adjoint */
  for (nadj=0; nadj<ts->numcost; nadj++) {
    KSPConvergedReason kspreason;
    ierr = KSPSolveTranspose(ksp,VecsSensiTemp[nadj],VecsDeltaLam[nadj]);CHKERRQ(ierr);
    ierr = KSPGetConvergedReason(ksp,&kspreason);CHKERRQ(ierr);
    if (kspreason < 0) {
      ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
      ierr = PetscInfo2(ts,"Step=%D, %Dth cost function, transposed linear solve fails, stopping 1st-order adjoint solve\n",ts->steps,nadj);CHKERRQ(ierr);
    }
  }

  /* Second-order adjoint */
  if (ts->vecs_sensi2) { /* U_{n+1} */
    if (!th->endpoint) SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_SUP,"Operation not implemented in TS_Theta");
    /* Get w1 at t_{n+1} from TLM matrix */
    ierr = MatDenseGetColumn(ts->mat_sensip,0,&xarr);CHKERRQ(ierr);
    ierr = VecPlaceArray(ts->vec_sensip_col,xarr);CHKERRQ(ierr);
    /* lambda_s^T F_UU w_1 */
    ierr = TSComputeIHessianProductFunctionUU(ts,th->stage_time,ts->vec_sol,VecsDeltaLam,ts->vec_sensip_col,ts->vecs_fuu);CHKERRQ(ierr);
    ierr = VecResetArray(ts->vec_sensip_col);CHKERRQ(ierr);
    ierr = MatDenseRestoreColumn(ts->mat_sensip,&xarr);CHKERRQ(ierr);
    /* lambda_s^T F_UP w_2 */
    ierr = TSComputeIHessianProductFunctionUP(ts,th->stage_time,ts->vec_sol,VecsDeltaLam,ts->vec_dir,ts->vecs_fup);CHKERRQ(ierr);
    for (nadj=0; nadj<ts->numcost; nadj++) { /* compute the residual */
      ierr = VecCopy(ts->vecs_sensi2[nadj],VecsSensi2Temp[nadj]);CHKERRQ(ierr);
      ierr = VecScale(VecsSensi2Temp[nadj],th->shift);CHKERRQ(ierr);
      ierr = VecAXPY(VecsSensi2Temp[nadj],-1.,ts->vecs_fuu[nadj]);CHKERRQ(ierr);
      if (ts->vecs_fup) {
        ierr = VecAXPY(VecsSensi2Temp[nadj],-1.,ts->vecs_fup[nadj]);CHKERRQ(ierr);
      }
    }
    /* Solve stage equation LHS X = RHS for second-order adjoint */
    for (nadj=0; nadj<ts->numcost; nadj++) {
      KSPConvergedReason kspreason;
      ierr = KSPSolveTranspose(ksp,VecsSensi2Temp[nadj],VecsDeltaLam2[nadj]);CHKERRQ(ierr);
      ierr = KSPGetConvergedReason(ksp,&kspreason);CHKERRQ(ierr);
      if (kspreason < 0) {
        ts->reason = TSADJOINT_DIVERGED_LINEAR_SOLVE;
        ierr = PetscInfo2(ts,"Step=%D, %Dth cost function, transposed linear solve fails, stopping 2nd-order adjoint solve\n",ts->steps,nadj);CHKERRQ(ierr);
      }
    }
  }

  /* Update sensitivities, and evaluate integrals if there is any */
  if (th->endpoint) { /* two-stage Theta methods with th->Theta!=1, th->Theta==1 leads to BEuler */
    th->shift      = 1./((th->Theta-1.)*adjoint_time_step);
    th->stage_time = adjoint_ptime;
    ierr           = TSComputeSNESJacobian(ts,th->X0,J,Jpre);CHKERRQ(ierr);
    ierr           = KSPSetOperators(ksp,J,Jpre);CHKERRQ(ierr);
    /* R_U at t_n */
    if (quadts) {
      ierr = TSComputeRHSJacobian(quadts,adjoint_ptime,th->X0,quadJ,NULL);CHKERRQ(ierr);
    }
    for (nadj=0; nadj<ts->numcost; nadj++) {
      ierr = MatMultTranspose(J,VecsDeltaLam[nadj],ts->vecs_sensi[nadj]);CHKERRQ(ierr);
      if (quadJ) {
        ierr = MatDenseGetColumn(quadJ,nadj,&xarr);CHKERRQ(ierr);
        ierr = VecPlaceArray(ts->vec_drdu_col,xarr);CHKERRQ(ierr);
        ierr = VecAXPY(ts->vecs_sensi[nadj],-1.,ts->vec_drdu_col);CHKERRQ(ierr);
        ierr = VecResetArray(ts->vec_drdu_col);CHKERRQ(ierr);
        ierr = MatDenseRestoreColumn(quadJ,&xarr);CHKERRQ(ierr);
      }
      ierr = VecScale(ts->vecs_sensi[nadj],1./th->shift);CHKERRQ(ierr);
    }

    /* Second-order adjoint */
    if (ts->vecs_sensi2) { /* U_n */
      /* Get w1 at t_n from TLM matrix */
      ierr = MatDenseGetColumn(th->MatFwdSensip0,0,&xarr);CHKERRQ(ierr);
      ierr = VecPlaceArray(ts->vec_sensip_col,xarr);CHKERRQ(ierr);
      /* lambda_s^T F_UU w_1 */
      ierr = TSComputeIHessianProductFunctionUU(ts,adjoint_ptime,th->X0,VecsDeltaLam,ts->vec_sensip_col,ts->vecs_fuu);CHKERRQ(ierr);
      ierr = VecResetArray(ts->vec_sensip_col);CHKERRQ(ierr);
      ierr = MatDenseRestoreColumn(th->MatFwdSensip0,&xarr);CHKERRQ(ierr);
      /* lambda_s^T F_UU w_2 */
      ierr = TSComputeIHessianProductFunctionUP(ts,adjoint_ptime,th->X0,VecsDeltaLam,ts->vec_dir,ts->vecs_fup);CHKERRQ(ierr);
      for (nadj=0; nadj<ts->numcost; nadj++) {
        /* 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  */
        ierr = MatMultTranspose(J,VecsDeltaLam2[nadj],ts->vecs_sensi2[nadj]);CHKERRQ(ierr);
        ierr = VecAXPY(ts->vecs_sensi2[nadj],1.,ts->vecs_fuu[nadj]);CHKERRQ(ierr);
        if (ts->vecs_fup) {
          ierr = VecAXPY(ts->vecs_sensi2[nadj],1.,ts->vecs_fup[nadj]);CHKERRQ(ierr);
        }
        ierr = VecScale(ts->vecs_sensi2[nadj],1./th->shift);CHKERRQ(ierr);
      }
    }

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

    if (ts->vecs_sensip) { /* sensitivities wrt parameters */
      /* U_{n+1} */
      th->shift = 1.0/(adjoint_time_step*th->Theta);
      ierr = VecAXPBYPCZ(th->Xdot,-th->shift,th->shift,0,th->X0,ts->vec_sol);CHKERRQ(ierr);
      ierr = TSComputeIJacobianP(ts,th->stage_time,ts->vec_sol,th->Xdot,-1./(th->Theta*adjoint_time_step),ts->Jacp,PETSC_FALSE);CHKERRQ(ierr);
      if (quadts) {
        ierr = TSComputeRHSJacobianP(quadts,th->stage_time,ts->vec_sol,quadJp);CHKERRQ(ierr);
      }
      for (nadj=0; nadj<ts->numcost; nadj++) {
        ierr = MatMultTranspose(ts->Jacp,VecsDeltaLam[nadj],VecsDeltaMu[nadj]);CHKERRQ(ierr);
        ierr = VecAXPY(ts->vecs_sensip[nadj],-adjoint_time_step*th->Theta,VecsDeltaMu[nadj]);CHKERRQ(ierr);
        if (quadJp) {
          ierr = MatDenseGetColumn(quadJp,nadj,&xarr);CHKERRQ(ierr);
          ierr = VecPlaceArray(ts->vec_drdp_col,xarr);CHKERRQ(ierr);
          ierr = VecAXPY(ts->vecs_sensip[nadj],adjoint_time_step*th->Theta,ts->vec_drdp_col);CHKERRQ(ierr);
          ierr = VecResetArray(ts->vec_drdp_col);CHKERRQ(ierr);
          ierr = MatDenseRestoreColumn(quadJp,&xarr);CHKERRQ(ierr);
        }
      }
      if (ts->vecs_sensi2p) { /* second-order */
        /* Get w1 at t_{n+1} from TLM matrix */
        ierr = MatDenseGetColumn(ts->mat_sensip,0,&xarr);CHKERRQ(ierr);
        ierr = VecPlaceArray(ts->vec_sensip_col,xarr);CHKERRQ(ierr);
        /* lambda_s^T F_PU w_1 */
        ierr = TSComputeIHessianProductFunctionPU(ts,th->stage_time,ts->vec_sol,VecsDeltaLam,ts->vec_sensip_col,ts->vecs_fpu);CHKERRQ(ierr);
        ierr = VecResetArray(ts->vec_sensip_col);CHKERRQ(ierr);
        ierr = MatDenseRestoreColumn(ts->mat_sensip,&xarr);CHKERRQ(ierr);

        /* lambda_s^T F_PP w_2 */
        ierr = TSComputeIHessianProductFunctionPP(ts,th->stage_time,ts->vec_sol,VecsDeltaLam,ts->vec_dir,ts->vecs_fpp);CHKERRQ(ierr);
        for (nadj=0; nadj<ts->numcost; nadj++) {
          /* 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)  */
          ierr = MatMultTranspose(ts->Jacp,VecsDeltaLam2[nadj],VecsDeltaMu2[nadj]);CHKERRQ(ierr);
          ierr = VecAXPY(ts->vecs_sensi2p[nadj],-adjoint_time_step*th->Theta,VecsDeltaMu2[nadj]);CHKERRQ(ierr);
          if (ts->vecs_fpu) {
            ierr = VecAXPY(ts->vecs_sensi2p[nadj],-adjoint_time_step*th->Theta,ts->vecs_fpu[nadj]);CHKERRQ(ierr);
          }
          if (ts->vecs_fpp) {
            ierr = VecAXPY(ts->vecs_sensi2p[nadj],-adjoint_time_step*th->Theta,ts->vecs_fpp[nadj]);CHKERRQ(ierr);
          }
        }
      }

      /* U_s */
      ierr = VecZeroEntries(th->Xdot);CHKERRQ(ierr);
      ierr = TSComputeIJacobianP(ts,adjoint_ptime,th->X0,th->Xdot,1./((th->Theta-1.0)*adjoint_time_step),ts->Jacp,PETSC_FALSE);CHKERRQ(ierr);
      if (quadts) {
        ierr = TSComputeRHSJacobianP(quadts,adjoint_ptime,th->X0,quadJp);CHKERRQ(ierr);
      }
      for (nadj=0; nadj<ts->numcost; nadj++) {
        ierr = MatMultTranspose(ts->Jacp,VecsDeltaLam[nadj],VecsDeltaMu[nadj]);CHKERRQ(ierr);
        ierr = VecAXPY(ts->vecs_sensip[nadj],-adjoint_time_step*(1.0-th->Theta),VecsDeltaMu[nadj]);CHKERRQ(ierr);
        if (quadJp) {
          ierr = MatDenseGetColumn(quadJp,nadj,&xarr);CHKERRQ(ierr);
          ierr = VecPlaceArray(ts->vec_drdp_col,xarr);CHKERRQ(ierr);
          ierr = VecAXPY(ts->vecs_sensip[nadj],adjoint_time_step*(1.0-th->Theta),ts->vec_drdp_col);CHKERRQ(ierr);
          ierr = VecResetArray(ts->vec_drdp_col);CHKERRQ(ierr);
          ierr = MatDenseRestoreColumn(quadJp,&xarr);CHKERRQ(ierr);
        }
        if (ts->vecs_sensi2p) { /* second-order */
          /* Get w1 at t_n from TLM matrix */
          ierr = MatDenseGetColumn(th->MatFwdSensip0,0,&xarr);CHKERRQ(ierr);
          ierr = VecPlaceArray(ts->vec_sensip_col,xarr);CHKERRQ(ierr);
          /* lambda_s^T F_PU w_1 */
          ierr = TSComputeIHessianProductFunctionPU(ts,adjoint_ptime,th->X0,VecsDeltaLam,ts->vec_sensip_col,ts->vecs_fpu);CHKERRQ(ierr);
          ierr = VecResetArray(ts->vec_sensip_col);CHKERRQ(ierr);
          ierr = MatDenseRestoreColumn(th->MatFwdSensip0,&xarr);CHKERRQ(ierr);
          /* lambda_s^T F_PP w_2 */
          ierr = TSComputeIHessianProductFunctionPP(ts,adjoint_ptime,th->X0,VecsDeltaLam,ts->vec_dir,ts->vecs_fpp);CHKERRQ(ierr);
          for (nadj=0; nadj<ts->numcost; nadj++) {
            /* 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) */
            ierr = MatMultTranspose(ts->Jacp,VecsDeltaLam2[nadj],VecsDeltaMu2[nadj]);CHKERRQ(ierr);
            ierr = VecAXPY(ts->vecs_sensi2p[nadj],-adjoint_time_step*(1.0-th->Theta),VecsDeltaMu2[nadj]);CHKERRQ(ierr);
            if (ts->vecs_fpu) {
              ierr = VecAXPY(ts->vecs_sensi2p[nadj],-adjoint_time_step*(1.0-th->Theta),ts->vecs_fpu[nadj]);CHKERRQ(ierr);
            }
            if (ts->vecs_fpp) {
              ierr = VecAXPY(ts->vecs_sensi2p[nadj],-adjoint_time_step*(1.0-th->Theta),ts->vecs_fpp[nadj]);CHKERRQ(ierr);
            }
          }
        }
      }
    }
  } else { /* one-stage case */
    th->shift = 0.0;
    ierr      = TSComputeSNESJacobian(ts,th->X,J,Jpre);CHKERRQ(ierr); /* get -f_y */
    ierr      = KSPSetOperators(ksp,J,Jpre);CHKERRQ(ierr);
    if (quadts) {
      ierr  = TSComputeRHSJacobian(quadts,th->stage_time,th->X,quadJ,NULL);CHKERRQ(ierr);
    }
    for (nadj=0; nadj<ts->numcost; nadj++) {
      ierr = MatMultTranspose(J,VecsDeltaLam[nadj],VecsSensiTemp[nadj]);CHKERRQ(ierr);
      ierr = VecAXPY(ts->vecs_sensi[nadj],-adjoint_time_step,VecsSensiTemp[nadj]);CHKERRQ(ierr);
      if (quadJ) {
        ierr = MatDenseGetColumn(quadJ,nadj,&xarr);CHKERRQ(ierr);
        ierr = VecPlaceArray(ts->vec_drdu_col,xarr);CHKERRQ(ierr);
        ierr = VecAXPY(ts->vecs_sensi[nadj],adjoint_time_step,ts->vec_drdu_col);CHKERRQ(ierr);
        ierr = VecResetArray(ts->vec_drdu_col);CHKERRQ(ierr);
        ierr = MatDenseRestoreColumn(quadJ,&xarr);CHKERRQ(ierr);
      }
    }
    if (ts->vecs_sensip) {
      th->shift = 1.0/(adjoint_time_step*th->Theta);
      ierr = VecAXPBYPCZ(th->Xdot,-th->shift,th->shift,0,th->X0,th->X);CHKERRQ(ierr);
      ierr = TSComputeIJacobianP(ts,th->stage_time,th->X,th->Xdot,th->shift,ts->Jacp,PETSC_FALSE);CHKERRQ(ierr);
      if (quadts) {
        ierr = TSComputeRHSJacobianP(quadts,th->stage_time,th->X,quadJp);CHKERRQ(ierr);
      }
      for (nadj=0; nadj<ts->numcost; nadj++) {
        ierr = MatMultTranspose(ts->Jacp,VecsDeltaLam[nadj],VecsDeltaMu[nadj]);CHKERRQ(ierr);
        ierr = VecAXPY(ts->vecs_sensip[nadj],-adjoint_time_step,VecsDeltaMu[nadj]);CHKERRQ(ierr);
        if (quadJp) {
          ierr = MatDenseGetColumn(quadJp,nadj,&xarr);CHKERRQ(ierr);
          ierr = VecPlaceArray(ts->vec_drdp_col,xarr);CHKERRQ(ierr);
          ierr = VecAXPY(ts->vecs_sensip[nadj],adjoint_time_step,ts->vec_drdp_col);CHKERRQ(ierr);
          ierr = VecResetArray(ts->vec_drdp_col);CHKERRQ(ierr);
          ierr = MatDenseRestoreColumn(quadJp,&xarr);CHKERRQ(ierr);
        }
      }
    }
  }

  th->status = TS_STEP_COMPLETE;
  PetscFunctionReturn(0);
}

static PetscErrorCode TSInterpolate_Theta(TS ts,PetscReal t,Vec X)
{
  TS_Theta       *th = (TS_Theta*)ts->data;
  PetscReal      dt  = t - ts->ptime;
  PetscErrorCode ierr;

  PetscFunctionBegin;
  ierr = VecCopy(ts->vec_sol,th->X);CHKERRQ(ierr);
  if (th->endpoint) dt *= th->Theta;
  ierr = VecWAXPY(X,dt,th->Xdot,th->X);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

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

  PetscFunctionBegin;
  if (!th->vec_sol_prev) {*wlte = -1; PetscFunctionReturn(0);}
  /* Cannot compute LTE in first step or in restart after event */
  if (ts->steprestart) {*wlte = -1; PetscFunctionReturn(0);}
  /* Compute LTE using backward differences with non-constant time step */
  {
    PetscReal   h = ts->time_step, h_prev = ts->ptime - ts->ptime_prev;
    PetscReal   a = 1 + h_prev/h;
    PetscScalar scal[3]; Vec vecs[3];
    scal[0] = +1/a; scal[1] = -1/(a-1); scal[2] = +1/(a*(a-1));
    vecs[0] = X;    vecs[1] = th->X0;   vecs[2] = th->vec_sol_prev;
    ierr = VecCopy(X,Y);CHKERRQ(ierr);
    ierr = VecMAXPY(Y,3,scal,vecs);CHKERRQ(ierr);
    ierr = TSErrorWeightedNorm(ts,X,Y,wnormtype,wlte,&wltea,&wlter);CHKERRQ(ierr);
  }
  if (order) *order = 2;
  PetscFunctionReturn(0);
}

static PetscErrorCode TSRollBack_Theta(TS ts)
{
  TS_Theta       *th = (TS_Theta*)ts->data;
  TS             quadts = ts->quadraturets;
  PetscErrorCode ierr;

  PetscFunctionBegin;
  ierr = VecCopy(th->X0,ts->vec_sol);CHKERRQ(ierr);
  if (quadts && ts->costintegralfwd) {
    ierr = VecCopy(th->VecCostIntegral0,quadts->vec_sol);CHKERRQ(ierr);
  }
  th->status = TS_STEP_INCOMPLETE;
  if (ts->mat_sensip) {
    ierr = MatCopy(th->MatFwdSensip0,ts->mat_sensip,SAME_NONZERO_PATTERN);CHKERRQ(ierr);
  }
  if (quadts && quadts->mat_sensip) {
    ierr = MatCopy(th->MatIntegralSensip0,quadts->mat_sensip,SAME_NONZERO_PATTERN);CHKERRQ(ierr);
  }
  PetscFunctionReturn(0);
}

static PetscErrorCode TSForwardStep_Theta(TS ts)
{
  TS_Theta       *th = (TS_Theta*)ts->data;
  TS             quadts = ts->quadraturets;
  Mat            MatDeltaFwdSensip = th->MatDeltaFwdSensip;
  Vec            VecDeltaFwdSensipCol = th->VecDeltaFwdSensipCol;
  PetscInt       ntlm;
  KSP            ksp;
  Mat            J,Jpre,quadJ = NULL,quadJp = NULL;
  PetscScalar    *barr,*xarr;
  PetscReal      previous_shift;
  PetscErrorCode ierr;

  PetscFunctionBegin;
  previous_shift = th->shift;
  ierr = MatCopy(ts->mat_sensip,th->MatFwdSensip0,SAME_NONZERO_PATTERN);CHKERRQ(ierr);

  if (quadts && quadts->mat_sensip) {
    ierr = MatCopy(quadts->mat_sensip,th->MatIntegralSensip0,SAME_NONZERO_PATTERN);CHKERRQ(ierr);
  }
  ierr = SNESGetKSP(ts->snes,&ksp);CHKERRQ(ierr);
  ierr = TSGetIJacobian(ts,&J,&Jpre,NULL,NULL);CHKERRQ(ierr);
  if (quadts) {
    ierr = TSGetRHSJacobian(quadts,&quadJ,NULL,NULL,NULL);CHKERRQ(ierr);
    ierr = TSGetRHSJacobianP(quadts,&quadJp,NULL,NULL);CHKERRQ(ierr);
  }

  /* Build RHS */
  if (th->endpoint) { /* 2-stage method*/
    th->shift = 1./((th->Theta-1.)*th->time_step0);
    ierr = TSComputeIJacobian(ts,th->ptime0,th->X0,th->Xdot,th->shift,J,Jpre,PETSC_FALSE);CHKERRQ(ierr);
    ierr = MatMatMult(J,ts->mat_sensip,MAT_REUSE_MATRIX,PETSC_DEFAULT,&MatDeltaFwdSensip);CHKERRQ(ierr);
    ierr = MatScale(MatDeltaFwdSensip,(th->Theta-1.)/th->Theta);CHKERRQ(ierr);

    /* Add the f_p forcing terms */
    if (ts->Jacp) {
      ierr = VecZeroEntries(th->Xdot);CHKERRQ(ierr);
      ierr = TSComputeIJacobianP(ts,th->ptime0,th->X0,th->Xdot,th->shift,ts->Jacp,PETSC_FALSE);CHKERRQ(ierr);
      ierr = MatAXPY(MatDeltaFwdSensip,(th->Theta-1.)/th->Theta,ts->Jacp,SUBSET_NONZERO_PATTERN);CHKERRQ(ierr);
      th->shift = previous_shift;
      ierr = VecAXPBYPCZ(th->Xdot,-th->shift,th->shift,0,th->X0,ts->vec_sol);CHKERRQ(ierr);
      ierr = TSComputeIJacobianP(ts,th->stage_time,ts->vec_sol,th->Xdot,th->shift,ts->Jacp,PETSC_FALSE);CHKERRQ(ierr);
      ierr = MatAXPY(MatDeltaFwdSensip,-1.,ts->Jacp,SUBSET_NONZERO_PATTERN);CHKERRQ(ierr);
    }
  } else { /* 1-stage method */
    th->shift = 0.0;
    ierr = TSComputeIJacobian(ts,th->stage_time,th->X,th->Xdot,th->shift,J,Jpre,PETSC_FALSE);CHKERRQ(ierr);
    ierr = MatMatMult(J,ts->mat_sensip,MAT_REUSE_MATRIX,PETSC_DEFAULT,&MatDeltaFwdSensip);CHKERRQ(ierr);
    ierr = MatScale(MatDeltaFwdSensip,-1.);CHKERRQ(ierr);

    /* Add the f_p forcing terms */
    if (ts->Jacp) {
      th->shift = previous_shift;
      ierr = VecAXPBYPCZ(th->Xdot,-th->shift,th->shift,0,th->X0,th->X);CHKERRQ(ierr);
      ierr = TSComputeIJacobianP(ts,th->stage_time,th->X,th->Xdot,th->shift,ts->Jacp,PETSC_FALSE);CHKERRQ(ierr);
      ierr = MatAXPY(MatDeltaFwdSensip,-1.,ts->Jacp,SUBSET_NONZERO_PATTERN);CHKERRQ(ierr);
    }
  }

  /* Build LHS */
  th->shift = previous_shift; /* recover the previous shift used in TSStep_Theta() */
  if (th->endpoint) {
    ierr = TSComputeIJacobian(ts,th->stage_time,ts->vec_sol,th->Xdot,th->shift,J,Jpre,PETSC_FALSE);CHKERRQ(ierr);
  } else {
    ierr = TSComputeIJacobian(ts,th->stage_time,th->X,th->Xdot,th->shift,J,Jpre,PETSC_FALSE);CHKERRQ(ierr);
  }
  ierr = KSPSetOperators(ksp,J,Jpre);CHKERRQ(ierr);

  /*
    Evaluate the first stage of integral gradients with the 2-stage method:
    drdu|t_n*S(t_n) + drdp|t_n
    This is done before the linear solve because the sensitivity variable S(t_n) will be propagated to S(t_{n+1})
  */
  if (th->endpoint) { /* 2-stage method only */
    if (quadts && quadts->mat_sensip) {
      ierr = TSComputeRHSJacobian(quadts,th->ptime0,th->X0,quadJ,NULL);CHKERRQ(ierr);
      ierr = TSComputeRHSJacobianP(quadts,th->ptime0,th->X0,quadJp);CHKERRQ(ierr);
      ierr = MatTransposeMatMult(ts->mat_sensip,quadJ,MAT_REUSE_MATRIX,PETSC_DEFAULT,&th->MatIntegralSensipTemp);CHKERRQ(ierr);
      ierr = MatAXPY(th->MatIntegralSensipTemp,1,quadJp,SAME_NONZERO_PATTERN);CHKERRQ(ierr);
      ierr = MatAXPY(quadts->mat_sensip,th->time_step0*(1.-th->Theta),th->MatIntegralSensipTemp,SAME_NONZERO_PATTERN);CHKERRQ(ierr);
    }
  }

  /* Solve the tangent linear equation for forward sensitivities to parameters */
  for (ntlm=0; ntlm<th->num_tlm; ntlm++) {
    KSPConvergedReason kspreason;
    ierr = MatDenseGetColumn(MatDeltaFwdSensip,ntlm,&barr);CHKERRQ(ierr);
    ierr = VecPlaceArray(VecDeltaFwdSensipCol,barr);CHKERRQ(ierr);
    if (th->endpoint) {
      ierr = MatDenseGetColumn(ts->mat_sensip,ntlm,&xarr);CHKERRQ(ierr);
      ierr = VecPlaceArray(ts->vec_sensip_col,xarr);CHKERRQ(ierr);
      ierr = KSPSolve(ksp,VecDeltaFwdSensipCol,ts->vec_sensip_col);CHKERRQ(ierr);
      ierr = VecResetArray(ts->vec_sensip_col);CHKERRQ(ierr);
      ierr = MatDenseRestoreColumn(ts->mat_sensip,&xarr);CHKERRQ(ierr);
    } else {
      ierr = KSPSolve(ksp,VecDeltaFwdSensipCol,VecDeltaFwdSensipCol);CHKERRQ(ierr);
    }
    ierr = KSPGetConvergedReason(ksp,&kspreason);CHKERRQ(ierr);
    if (kspreason < 0) {
      ts->reason = TSFORWARD_DIVERGED_LINEAR_SOLVE;
      ierr = PetscInfo2(ts,"Step=%D, %Dth tangent linear solve, linear solve fails, stopping tangent linear solve\n",ts->steps,ntlm);CHKERRQ(ierr);
    }
    ierr = VecResetArray(VecDeltaFwdSensipCol);CHKERRQ(ierr);
    ierr = MatDenseRestoreColumn(MatDeltaFwdSensip,&barr);CHKERRQ(ierr);
  }

  /*
    Evaluate the second stage of integral gradients with the 2-stage method:
    drdu|t_{n+1}*S(t_{n+1}) + drdp|t_{n+1}
  */
  if (quadts && quadts->mat_sensip) {
    if (!th->endpoint) {
      ierr = MatAXPY(ts->mat_sensip,1,MatDeltaFwdSensip,SAME_NONZERO_PATTERN);CHKERRQ(ierr); /* stage sensitivity */
      ierr = TSComputeRHSJacobian(quadts,th->stage_time,th->X,quadJ,NULL);CHKERRQ(ierr);
      ierr = TSComputeRHSJacobianP(quadts,th->stage_time,th->X,quadJp);CHKERRQ(ierr);
      ierr = MatTransposeMatMult(ts->mat_sensip,quadJ,MAT_REUSE_MATRIX,PETSC_DEFAULT,&th->MatIntegralSensipTemp);CHKERRQ(ierr);
      ierr = MatAXPY(th->MatIntegralSensipTemp,1,quadJp,SAME_NONZERO_PATTERN);CHKERRQ(ierr);
      ierr = MatAXPY(quadts->mat_sensip,th->time_step0,th->MatIntegralSensipTemp,SAME_NONZERO_PATTERN);CHKERRQ(ierr);
      ierr = MatAXPY(ts->mat_sensip,(1.-th->Theta)/th->Theta,MatDeltaFwdSensip,SAME_NONZERO_PATTERN);CHKERRQ(ierr);
    } else {
      ierr = TSComputeRHSJacobian(quadts,th->stage_time,ts->vec_sol,quadJ,NULL);CHKERRQ(ierr);
      ierr = TSComputeRHSJacobianP(quadts,th->stage_time,ts->vec_sol,quadJp);CHKERRQ(ierr);
      ierr = MatTransposeMatMult(ts->mat_sensip,quadJ,MAT_REUSE_MATRIX,PETSC_DEFAULT,&th->MatIntegralSensipTemp);CHKERRQ(ierr);
      ierr = MatAXPY(th->MatIntegralSensipTemp,1,quadJp,SAME_NONZERO_PATTERN);CHKERRQ(ierr);
      ierr = MatAXPY(quadts->mat_sensip,th->time_step0*th->Theta,th->MatIntegralSensipTemp,SAME_NONZERO_PATTERN);CHKERRQ(ierr);
    }
  } else {
    if (!th->endpoint) {
      ierr = MatAXPY(ts->mat_sensip,1./th->Theta,MatDeltaFwdSensip,SAME_NONZERO_PATTERN);CHKERRQ(ierr);
    }
  }
  PetscFunctionReturn(0);
}

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

  PetscFunctionBegin;
  if (ns) {
    if (!th->endpoint && th->Theta != 1.0) *ns = 1; /* midpoint form */
    else *ns = 2; /* endpoint form */
  }
  if (stagesensip) {
    if (!th->endpoint && th->Theta != 1.0) {
      th->MatFwdStages[0] = th->MatDeltaFwdSensip;
    } else {
      th->MatFwdStages[0] = th->MatFwdSensip0;
      th->MatFwdStages[1] = ts->mat_sensip; /* stiffly accurate */
    }
    *stagesensip = th->MatFwdStages;
  }
  PetscFunctionReturn(0);
}

/*------------------------------------------------------------*/
static PetscErrorCode TSReset_Theta(TS ts)
{
  TS_Theta       *th = (TS_Theta*)ts->data;
  PetscErrorCode ierr;

  PetscFunctionBegin;
  ierr = VecDestroy(&th->X);CHKERRQ(ierr);
  ierr = VecDestroy(&th->Xdot);CHKERRQ(ierr);
  ierr = VecDestroy(&th->X0);CHKERRQ(ierr);
  ierr = VecDestroy(&th->affine);CHKERRQ(ierr);

  ierr = VecDestroy(&th->vec_sol_prev);CHKERRQ(ierr);
  ierr = VecDestroy(&th->vec_lte_work);CHKERRQ(ierr);

  ierr = VecDestroy(&th->VecCostIntegral0);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

static PetscErrorCode TSAdjointReset_Theta(TS ts)
{
  TS_Theta       *th = (TS_Theta*)ts->data;
  PetscErrorCode ierr;

  PetscFunctionBegin;
  ierr = VecDestroyVecs(ts->numcost,&th->VecsDeltaLam);CHKERRQ(ierr);
  ierr = VecDestroyVecs(ts->numcost,&th->VecsDeltaMu);CHKERRQ(ierr);
  ierr = VecDestroyVecs(ts->numcost,&th->VecsDeltaLam2);CHKERRQ(ierr);
  ierr = VecDestroyVecs(ts->numcost,&th->VecsDeltaMu2);CHKERRQ(ierr);
  ierr = VecDestroyVecs(ts->numcost,&th->VecsSensiTemp);CHKERRQ(ierr);
  ierr = VecDestroyVecs(ts->numcost,&th->VecsSensi2Temp);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

static PetscErrorCode TSDestroy_Theta(TS ts)
{
  PetscErrorCode ierr;

  PetscFunctionBegin;
  ierr = TSReset_Theta(ts);CHKERRQ(ierr);
  if (ts->dm) {
    ierr = DMCoarsenHookRemove(ts->dm,DMCoarsenHook_TSTheta,DMRestrictHook_TSTheta,ts);CHKERRQ(ierr);
    ierr = DMSubDomainHookRemove(ts->dm,DMSubDomainHook_TSTheta,DMSubDomainRestrictHook_TSTheta,ts);CHKERRQ(ierr);
  }
  ierr = PetscFree(ts->data);CHKERRQ(ierr);
  ierr = PetscObjectComposeFunction((PetscObject)ts,"TSThetaGetTheta_C",NULL);CHKERRQ(ierr);
  ierr = PetscObjectComposeFunction((PetscObject)ts,"TSThetaSetTheta_C",NULL);CHKERRQ(ierr);
  ierr = PetscObjectComposeFunction((PetscObject)ts,"TSThetaGetEndpoint_C",NULL);CHKERRQ(ierr);
  ierr = PetscObjectComposeFunction((PetscObject)ts,"TSThetaSetEndpoint_C",NULL);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

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

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

  PetscFunctionBegin;
  ierr = SNESGetDM(snes,&dm);CHKERRQ(ierr);
  /* When using the endpoint variant, this is actually 1/Theta * Xdot */
  ierr = TSThetaGetX0AndXdot(ts,dm,&X0,&Xdot);CHKERRQ(ierr);
  if (x != X0) {
    ierr = VecAXPBYPCZ(Xdot,-shift,shift,0,X0,x);CHKERRQ(ierr);
  } else {
    ierr = VecZeroEntries(Xdot);CHKERRQ(ierr);
  }
  /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */
  dmsave = ts->dm;
  ts->dm = dm;
  ierr   = TSComputeIFunction(ts,th->stage_time,x,Xdot,y,PETSC_FALSE);CHKERRQ(ierr);
  ts->dm = dmsave;
  ierr   = TSThetaRestoreX0AndXdot(ts,dm,&X0,&Xdot);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

static PetscErrorCode SNESTSFormJacobian_Theta(SNES snes,Vec x,Mat A,Mat B,TS ts)
{
  TS_Theta       *th = (TS_Theta*)ts->data;
  PetscErrorCode ierr;
  Vec            Xdot;
  DM             dm,dmsave;
  PetscReal      shift = th->shift;

  PetscFunctionBegin;
  ierr = SNESGetDM(snes,&dm);CHKERRQ(ierr);
  /* Xdot has already been computed in SNESTSFormFunction_Theta (SNES guarantees this) */
  ierr = TSThetaGetX0AndXdot(ts,dm,NULL,&Xdot);CHKERRQ(ierr);

  dmsave = ts->dm;
  ts->dm = dm;
  ierr   = TSComputeIJacobian(ts,th->stage_time,x,Xdot,shift,A,B,PETSC_FALSE);CHKERRQ(ierr);
  ts->dm = dmsave;
  ierr   = TSThetaRestoreX0AndXdot(ts,dm,NULL,&Xdot);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

static PetscErrorCode TSForwardSetUp_Theta(TS ts)
{
  TS_Theta       *th = (TS_Theta*)ts->data;
  TS             quadts = ts->quadraturets;
  PetscErrorCode ierr;

  PetscFunctionBegin;
  /* combine sensitivities to parameters and sensitivities to initial values into one array */
  th->num_tlm = ts->num_parameters;
  ierr = MatDuplicate(ts->mat_sensip,MAT_DO_NOT_COPY_VALUES,&th->MatDeltaFwdSensip);CHKERRQ(ierr);
  if (quadts && quadts->mat_sensip) {
    ierr = MatDuplicate(quadts->mat_sensip,MAT_DO_NOT_COPY_VALUES,&th->MatIntegralSensipTemp);CHKERRQ(ierr);
    ierr = MatDuplicate(quadts->mat_sensip,MAT_DO_NOT_COPY_VALUES,&th->MatIntegralSensip0);CHKERRQ(ierr);
  }
  /* backup sensitivity results for roll-backs */
  ierr = MatDuplicate(ts->mat_sensip,MAT_DO_NOT_COPY_VALUES,&th->MatFwdSensip0);CHKERRQ(ierr);

  ierr = VecDuplicate(ts->vec_sol,&th->VecDeltaFwdSensipCol);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

static PetscErrorCode TSForwardReset_Theta(TS ts)
{
  TS_Theta       *th = (TS_Theta*)ts->data;
  TS             quadts = ts->quadraturets;
  PetscErrorCode ierr;

  PetscFunctionBegin;
  if (quadts && quadts->mat_sensip) {
    ierr = MatDestroy(&th->MatIntegralSensipTemp);CHKERRQ(ierr);
    ierr = MatDestroy(&th->MatIntegralSensip0);CHKERRQ(ierr);
  }
  ierr = VecDestroy(&th->VecDeltaFwdSensipCol);CHKERRQ(ierr);
  ierr = MatDestroy(&th->MatDeltaFwdSensip);CHKERRQ(ierr);
  ierr = MatDestroy(&th->MatFwdSensip0);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

static PetscErrorCode TSSetUp_Theta(TS ts)
{
  TS_Theta       *th = (TS_Theta*)ts->data;
  TS             quadts = ts->quadraturets;
  PetscBool      match;
  PetscErrorCode ierr;

  PetscFunctionBegin;
  if (!th->VecCostIntegral0 && quadts && ts->costintegralfwd) { /* back up cost integral */
    ierr = VecDuplicate(quadts->vec_sol,&th->VecCostIntegral0);CHKERRQ(ierr);
  }
  if (!th->X) {
    ierr = VecDuplicate(ts->vec_sol,&th->X);CHKERRQ(ierr);
  }
  if (!th->Xdot) {
    ierr = VecDuplicate(ts->vec_sol,&th->Xdot);CHKERRQ(ierr);
  }
  if (!th->X0) {
    ierr = VecDuplicate(ts->vec_sol,&th->X0);CHKERRQ(ierr);
  }
  if (th->endpoint) {
    ierr = VecDuplicate(ts->vec_sol,&th->affine);CHKERRQ(ierr);
  }

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

  ierr = TSGetDM(ts,&ts->dm);CHKERRQ(ierr);
  ierr = DMCoarsenHookAdd(ts->dm,DMCoarsenHook_TSTheta,DMRestrictHook_TSTheta,ts);CHKERRQ(ierr);
  ierr = DMSubDomainHookAdd(ts->dm,DMSubDomainHook_TSTheta,DMSubDomainRestrictHook_TSTheta,ts);CHKERRQ(ierr);

  ierr = TSGetAdapt(ts,&ts->adapt);CHKERRQ(ierr);
  ierr = TSAdaptCandidatesClear(ts->adapt);CHKERRQ(ierr);
  ierr = PetscObjectTypeCompare((PetscObject)ts->adapt,TSADAPTNONE,&match);CHKERRQ(ierr);
  if (!match) {
    ierr = VecDuplicate(ts->vec_sol,&th->vec_sol_prev);CHKERRQ(ierr);
    ierr = VecDuplicate(ts->vec_sol,&th->vec_lte_work);CHKERRQ(ierr);
  }
  ierr = TSGetSNES(ts,&ts->snes);CHKERRQ(ierr);

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

/*------------------------------------------------------------*/

static PetscErrorCode TSAdjointSetUp_Theta(TS ts)
{
  TS_Theta       *th = (TS_Theta*)ts->data;
  PetscErrorCode ierr;

  PetscFunctionBegin;
  ierr = VecDuplicateVecs(ts->vecs_sensi[0],ts->numcost,&th->VecsDeltaLam);CHKERRQ(ierr);
  ierr = VecDuplicateVecs(ts->vecs_sensi[0],ts->numcost,&th->VecsSensiTemp);CHKERRQ(ierr);
  if (ts->vecs_sensip) {
    ierr = VecDuplicateVecs(ts->vecs_sensip[0],ts->numcost,&th->VecsDeltaMu);CHKERRQ(ierr);
  }
  if (ts->vecs_sensi2) {
    ierr = VecDuplicateVecs(ts->vecs_sensi[0],ts->numcost,&th->VecsDeltaLam2);CHKERRQ(ierr);
    ierr = VecDuplicateVecs(ts->vecs_sensi2[0],ts->numcost,&th->VecsSensi2Temp);CHKERRQ(ierr);
    /* hack ts to make implicit TS solver work when users provide only explicit versions of callbacks (RHSFunction,RHSJacobian,RHSHessian etc.) */
    if (!ts->ihessianproduct_fuu) ts->vecs_fuu = ts->vecs_guu;
    if (!ts->ihessianproduct_fup) ts->vecs_fup = ts->vecs_gup;
  }
  if (ts->vecs_sensi2p) {
    ierr = VecDuplicateVecs(ts->vecs_sensi2p[0],ts->numcost,&th->VecsDeltaMu2);CHKERRQ(ierr);
    /* hack ts to make implicit TS solver work when users provide only explicit versions of callbacks (RHSFunction,RHSJacobian,RHSHessian etc.) */
    if (!ts->ihessianproduct_fpu) ts->vecs_fpu = ts->vecs_gpu;
    if (!ts->ihessianproduct_fpp) ts->vecs_fpp = ts->vecs_gpp;
  }
  PetscFunctionReturn(0);
}

static PetscErrorCode TSSetFromOptions_Theta(PetscOptionItems *PetscOptionsObject,TS ts)
{
  TS_Theta       *th = (TS_Theta*)ts->data;
  PetscErrorCode ierr;

  PetscFunctionBegin;
  ierr = PetscOptionsHead(PetscOptionsObject,"Theta ODE solver options");CHKERRQ(ierr);
  {
    ierr = PetscOptionsReal("-ts_theta_theta","Location of stage (0<Theta<=1)","TSThetaSetTheta",th->Theta,&th->Theta,NULL);CHKERRQ(ierr);
    ierr = PetscOptionsBool("-ts_theta_endpoint","Use the endpoint instead of midpoint form of the Theta method","TSThetaSetEndpoint",th->endpoint,&th->endpoint,NULL);CHKERRQ(ierr);
    ierr = PetscOptionsBool("-ts_theta_initial_guess_extrapolate","Extrapolate stage initial guess from previous solution (sometimes unstable)","TSThetaSetExtrapolate",th->extrapolate,&th->extrapolate,NULL);CHKERRQ(ierr);
  }
  ierr = PetscOptionsTail();CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

static PetscErrorCode TSView_Theta(TS ts,PetscViewer viewer)
{
  TS_Theta       *th = (TS_Theta*)ts->data;
  PetscBool      iascii;
  PetscErrorCode ierr;

  PetscFunctionBegin;
  ierr = PetscObjectTypeCompare((PetscObject)viewer,PETSCVIEWERASCII,&iascii);CHKERRQ(ierr);
  if (iascii) {
    ierr = PetscViewerASCIIPrintf(viewer,"  Theta=%g\n",(double)th->Theta);CHKERRQ(ierr);
    ierr = PetscViewerASCIIPrintf(viewer,"  Extrapolation=%s\n",th->extrapolate ? "yes" : "no");CHKERRQ(ierr);
  }
  PetscFunctionReturn(0);
}

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

  PetscFunctionBegin;
  *theta = th->Theta;
  PetscFunctionReturn(0);
}

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

  PetscFunctionBegin;
  if (theta <= 0 || 1 < theta) SETERRQ1(PetscObjectComm((PetscObject)ts),PETSC_ERR_ARG_OUTOFRANGE,"Theta %g not in range (0,1]",(double)theta);
  th->Theta = theta;
  th->order = (th->Theta == 0.5) ? 2 : 1;
  PetscFunctionReturn(0);
}

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

  PetscFunctionBegin;
  *endpoint = th->endpoint;
  PetscFunctionReturn(0);
}

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

  PetscFunctionBegin;
  th->endpoint = flg;
  PetscFunctionReturn(0);
}

#if defined(PETSC_HAVE_COMPLEX)
static PetscErrorCode TSComputeLinearStability_Theta(TS ts,PetscReal xr,PetscReal xi,PetscReal *yr,PetscReal *yi)
{
  PetscComplex   z   = xr + xi*PETSC_i,f;
  TS_Theta       *th = (TS_Theta*)ts->data;
  const PetscReal one = 1.0;

  PetscFunctionBegin;
  f   = (one + (one - th->Theta)*z)/(one - th->Theta*z);
  *yr = PetscRealPartComplex(f);
  *yi = PetscImaginaryPartComplex(f);
  PetscFunctionReturn(0);
}
#endif

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

  PetscFunctionBegin;
  if (ns) {
    if (!th->endpoint && th->Theta != 1.0) *ns = 1; /* midpoint form */
    else *ns = 2; /* endpoint form */
  }
  if (Y) {
    if (!th->endpoint && th->Theta != 1.0) {
      th->Stages[0] = th->X;
    } else {
      th->Stages[0] = th->X0;
      th->Stages[1] = ts->vec_sol; /* stiffly accurate */
    }
    *Y = th->Stages;
  }
  PetscFunctionReturn(0);
}

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

   Level: beginner

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

   Notes:
$  -ts_type theta -ts_theta_theta 1.0 corresponds to backward Euler (TSBEULER)
$  -ts_type theta -ts_theta_theta 0.5 corresponds to the implicit midpoint rule
$  -ts_type theta -ts_theta_theta 0.5 -ts_theta_endpoint corresponds to Crank-Nicholson (TSCN)

   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.

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

.vb
  Theta | Theta
  -------------
        |  1
.ve

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

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

.vb
  0 | 0         0
  1 | 1-Theta   Theta
  -------------------
    | 1-Theta   Theta
.ve

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

   To apply a diagonally implicit RK method to DAE, the stage formula

$  Y_i = X + h sum_j a_ij Y'_j

   is interpreted as a formula for Y'_i in terms of Y_i and known values (Y'_j, j<i)

.seealso:  TSCreate(), TS, TSSetType(), TSCN, TSBEULER, TSThetaSetTheta(), TSThetaSetEndpoint()

M*/
PETSC_EXTERN PetscErrorCode TSCreate_Theta(TS ts)
{
  TS_Theta       *th;
  PetscErrorCode ierr;

  PetscFunctionBegin;
  ts->ops->reset           = TSReset_Theta;
  ts->ops->adjointreset    = TSAdjointReset_Theta;
  ts->ops->destroy         = TSDestroy_Theta;
  ts->ops->view            = TSView_Theta;
  ts->ops->setup           = TSSetUp_Theta;
  ts->ops->adjointsetup    = TSAdjointSetUp_Theta;
  ts->ops->adjointreset    = TSAdjointReset_Theta;
  ts->ops->step            = TSStep_Theta;
  ts->ops->interpolate     = TSInterpolate_Theta;
  ts->ops->evaluatewlte    = TSEvaluateWLTE_Theta;
  ts->ops->rollback        = TSRollBack_Theta;
  ts->ops->setfromoptions  = TSSetFromOptions_Theta;
  ts->ops->snesfunction    = SNESTSFormFunction_Theta;
  ts->ops->snesjacobian    = SNESTSFormJacobian_Theta;
#if defined(PETSC_HAVE_COMPLEX)
  ts->ops->linearstability = TSComputeLinearStability_Theta;
#endif
  ts->ops->getstages       = TSGetStages_Theta;
  ts->ops->adjointstep     = TSAdjointStep_Theta;
  ts->ops->adjointintegral = TSAdjointCostIntegral_Theta;
  ts->ops->forwardintegral = TSForwardCostIntegral_Theta;
  ts->default_adapt_type   = TSADAPTNONE;

  ts->ops->forwardsetup     = TSForwardSetUp_Theta;
  ts->ops->forwardreset     = TSForwardReset_Theta;
  ts->ops->forwardstep      = TSForwardStep_Theta;
  ts->ops->forwardgetstages = TSForwardGetStages_Theta;

  ts->usessnes = PETSC_TRUE;

  ierr = PetscNewLog(ts,&th);CHKERRQ(ierr);
  ts->data = (void*)th;

  th->VecsDeltaLam    = NULL;
  th->VecsDeltaMu     = NULL;
  th->VecsSensiTemp   = NULL;
  th->VecsSensi2Temp  = NULL;

  th->extrapolate = PETSC_FALSE;
  th->Theta       = 0.5;
  th->order       = 2;
  ierr = PetscObjectComposeFunction((PetscObject)ts,"TSThetaGetTheta_C",TSThetaGetTheta_Theta);CHKERRQ(ierr);
  ierr = PetscObjectComposeFunction((PetscObject)ts,"TSThetaSetTheta_C",TSThetaSetTheta_Theta);CHKERRQ(ierr);
  ierr = PetscObjectComposeFunction((PetscObject)ts,"TSThetaGetEndpoint_C",TSThetaGetEndpoint_Theta);CHKERRQ(ierr);
  ierr = PetscObjectComposeFunction((PetscObject)ts,"TSThetaSetEndpoint_C",TSThetaSetEndpoint_Theta);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

/*@
  TSThetaGetTheta - Get the abscissa of the stage in (0,1].

  Not Collective

  Input Parameter:
.  ts - timestepping context

  Output Parameter:
.  theta - stage abscissa

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

  Level: Advanced

.seealso: TSThetaSetTheta()
@*/
PetscErrorCode  TSThetaGetTheta(TS ts,PetscReal *theta)
{
  PetscErrorCode ierr;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(ts,TS_CLASSID,1);
  PetscValidPointer(theta,2);
  ierr = PetscUseMethod(ts,"TSThetaGetTheta_C",(TS,PetscReal*),(ts,theta));CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

/*@
  TSThetaSetTheta - Set the abscissa of the stage in (0,1].

  Not Collective

  Input Parameter:
+  ts - timestepping context
-  theta - stage abscissa

  Options Database:
.  -ts_theta_theta <theta>

  Level: Intermediate

.seealso: TSThetaGetTheta()
@*/
PetscErrorCode  TSThetaSetTheta(TS ts,PetscReal theta)
{
  PetscErrorCode ierr;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(ts,TS_CLASSID,1);
  ierr = PetscTryMethod(ts,"TSThetaSetTheta_C",(TS,PetscReal),(ts,theta));CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

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

  Not Collective

  Input Parameter:
.  ts - timestepping context

  Output Parameter:
.  endpoint - PETSC_TRUE when using the endpoint variant

  Level: Advanced

.seealso: TSThetaSetEndpoint(), TSTHETA, TSCN
@*/
PetscErrorCode TSThetaGetEndpoint(TS ts,PetscBool *endpoint)
{
  PetscErrorCode ierr;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(ts,TS_CLASSID,1);
  PetscValidPointer(endpoint,2);
  ierr = PetscUseMethod(ts,"TSThetaGetEndpoint_C",(TS,PetscBool*),(ts,endpoint));CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

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

  Not Collective

  Input Parameter:
+  ts - timestepping context
-  flg - PETSC_TRUE to use the endpoint variant

  Options Database:
.  -ts_theta_endpoint <flg>

  Level: Intermediate

.seealso: TSTHETA, TSCN
@*/
PetscErrorCode TSThetaSetEndpoint(TS ts,PetscBool flg)
{
  PetscErrorCode ierr;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(ts,TS_CLASSID,1);
  ierr = PetscTryMethod(ts,"TSThetaSetEndpoint_C",(TS,PetscBool),(ts,flg));CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

/*
 * TSBEULER and TSCN are straightforward specializations of TSTHETA.
 * The creation functions for these specializations are below.
 */

static PetscErrorCode TSSetUp_BEuler(TS ts)
{
  TS_Theta       *th = (TS_Theta*)ts->data;
  PetscErrorCode ierr;

  PetscFunctionBegin;
  if (th->Theta != 1.0) SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_OPT_OVERWRITE,"Can not change the default value (1) of theta when using backward Euler\n");
  if (th->endpoint) SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_OPT_OVERWRITE,"Can not change to the endpoint form of the Theta methods when using backward Euler\n");
  ierr = TSSetUp_Theta(ts);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

static PetscErrorCode TSView_BEuler(TS ts,PetscViewer viewer)
{
  PetscFunctionBegin;
  PetscFunctionReturn(0);
}

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

  Level: beginner

  Notes:
  TSBEULER is equivalent to TSTHETA with Theta=1.0

$  -ts_type theta -ts_theta_theta 1.0

.seealso:  TSCreate(), TS, TSSetType(), TSEULER, TSCN, TSTHETA

M*/
PETSC_EXTERN PetscErrorCode TSCreate_BEuler(TS ts)
{
  PetscErrorCode ierr;

  PetscFunctionBegin;
  ierr = TSCreate_Theta(ts);CHKERRQ(ierr);
  ierr = TSThetaSetTheta(ts,1.0);CHKERRQ(ierr);
  ierr = TSThetaSetEndpoint(ts,PETSC_FALSE);CHKERRQ(ierr);
  ts->ops->setup = TSSetUp_BEuler;
  ts->ops->view  = TSView_BEuler;
  PetscFunctionReturn(0);
}

static PetscErrorCode TSSetUp_CN(TS ts)
{
  TS_Theta       *th = (TS_Theta*)ts->data;
  PetscErrorCode ierr;

  PetscFunctionBegin;
  if (th->Theta != 0.5) SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_OPT_OVERWRITE,"Can not change the default value (0.5) of theta when using Crank-Nicolson\n");
  if (!th->endpoint) SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_OPT_OVERWRITE,"Can not change to the midpoint form of the Theta methods when using Crank-Nicolson\n");
  ierr = TSSetUp_Theta(ts);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

static PetscErrorCode TSView_CN(TS ts,PetscViewer viewer)
{
  PetscFunctionBegin;
  PetscFunctionReturn(0);
}

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

  Level: beginner

  Notes:
  TSCN is equivalent to TSTHETA with Theta=0.5 and the "endpoint" option set. I.e.

$  -ts_type theta -ts_theta_theta 0.5 -ts_theta_endpoint

.seealso:  TSCreate(), TS, TSSetType(), TSBEULER, TSTHETA

M*/
PETSC_EXTERN PetscErrorCode TSCreate_CN(TS ts)
{
  PetscErrorCode ierr;

  PetscFunctionBegin;
  ierr = TSCreate_Theta(ts);CHKERRQ(ierr);
  ierr = TSThetaSetTheta(ts,0.5);CHKERRQ(ierr);
  ierr = TSThetaSetEndpoint(ts,PETSC_TRUE);CHKERRQ(ierr);
  ts->ops->setup = TSSetUp_CN;
  ts->ops->view  = TSView_CN;
  PetscFunctionReturn(0);
}