PetscErrorCode TSPrecond_Sundials(realtype tn,N_Vector y,N_Vector fy,booleantype jok,booleantype *jcurPtr,
realtype _gamma,void *P_data,N_Vector vtemp1,N_Vector vtemp2,N_Vector vtemp3)
{
TS ts = (TS) P_data;
TS_Sundials *cvode = (TS_Sundials*)ts->data;
PC pc;
PetscErrorCode ierr;
Mat J,P;
Vec yy = cvode->w1,yydot = cvode->ydot;
PetscReal gm = (PetscReal)_gamma;
MatStructure str = DIFFERENT_NONZERO_PATTERN;
PetscScalar *y_data;
PetscFunctionBegin;
ierr = TSGetIJacobian(ts,&J,&P,NULL,NULL);CHKERRQ(ierr);
y_data = (PetscScalar*) N_VGetArrayPointer(y);
ierr = VecPlaceArray(yy,y_data);CHKERRQ(ierr);
ierr = VecZeroEntries(yydot);CHKERRQ(ierr); /* The Jacobian is independent of Ydot for ODE which is all that CVode works for */
/* compute the shifted Jacobian (1/gm)*I + Jrest */
ierr = TSComputeIJacobian(ts,ts->ptime,yy,yydot,1/gm,&J,&P,&str,PETSC_FALSE);CHKERRQ(ierr);
ierr = VecResetArray(yy);CHKERRQ(ierr);
ierr = MatScale(P,gm);CHKERRQ(ierr); /* turn into I-gm*Jrest, J is not used by Sundials */
*jcurPtr = TRUE;
ierr = TSSundialsGetPC(ts,&pc);CHKERRQ(ierr);
ierr = PCSetOperators(pc,J,P,str);CHKERRQ(ierr);
PetscFunctionReturn(0);
}
static PetscErrorCode TSAdjointStep_Theta(TS ts)
{
TS_Theta *th = (TS_Theta*)ts->data;
Vec *VecsDeltaLam = th->VecsDeltaLam,*VecsDeltaMu = th->VecsDeltaMu,*VecsSensiTemp = th->VecsSensiTemp;
PetscInt nadj;
PetscErrorCode ierr;
Mat J,Jp;
KSP ksp;
PetscReal shift;
PetscFunctionBegin;
th->status = TS_STEP_INCOMPLETE;
ierr = SNESGetKSP(ts->snes,&ksp);CHKERRQ(ierr);
ierr = TSGetIJacobian(ts,&J,&Jp,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 = ts->ptime + (th->endpoint ? ts->time_step : (1.-th->Theta)*ts->time_step); /* time_step is negative*/
th->ptime = ts->ptime + ts->time_step;
/* Build RHS */
if (ts->vec_costintegral) { /* Cost function has an integral term */
if (th->endpoint) {
ierr = TSAdjointComputeDRDYFunction(ts,ts->ptime,ts->vec_sol,ts->vecs_drdy);CHKERRQ(ierr);
}else {
ierr = TSAdjointComputeDRDYFunction(ts,th->stage_time,th->X,ts->vecs_drdy);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*ts->time_step));CHKERRQ(ierr);
if (ts->vec_costintegral) {
ierr = VecAXPY(VecsSensiTemp[nadj],1.,ts->vecs_drdy[nadj]);CHKERRQ(ierr);
}
}
/* Build LHS */
shift = -1./(th->Theta*ts->time_step);
if (th->endpoint) {
ierr = TSComputeIJacobian(ts,ts->ptime,ts->vec_sol,th->Xdot,shift,J,Jp,PETSC_FALSE);CHKERRQ(ierr);
}else {
ierr = TSComputeIJacobian(ts,th->stage_time,th->X,th->Xdot,shift,J,Jp,PETSC_FALSE);CHKERRQ(ierr);
}
ierr = KSPSetOperators(ksp,J,Jp);CHKERRQ(ierr);
/* Solve LHS X = RHS */
for (nadj=0; nadj<ts->numcost; nadj++) {
ierr = KSPSolveTranspose(ksp,VecsSensiTemp[nadj],VecsDeltaLam[nadj]);CHKERRQ(ierr);
}
/* Update sensitivities, and evaluate integrals if there is any */
if(th->endpoint) { /* two-stage case */
if (th->Theta!=1.) {
shift = -1./((th->Theta-1.)*ts->time_step);
ierr = TSComputeIJacobian(ts,th->ptime,th->X0,th->Xdot,shift,J,Jp,PETSC_FALSE);CHKERRQ(ierr);
if (ts->vec_costintegral) {
ierr = TSAdjointComputeDRDYFunction(ts,th->ptime,th->X0,ts->vecs_drdy);CHKERRQ(ierr);
}
for (nadj=0; nadj<ts->numcost; nadj++) {
ierr = MatMultTranspose(J,VecsDeltaLam[nadj],ts->vecs_sensi[nadj]);CHKERRQ(ierr);
if (ts->vec_costintegral) {
ierr = VecAXPY(ts->vecs_sensi[nadj],-1.,ts->vecs_drdy[nadj]);CHKERRQ(ierr);
}
ierr = VecScale(ts->vecs_sensi[nadj],1./shift);CHKERRQ(ierr);
}
}else { /* backward Euler */
shift = 0.0;
ierr = TSComputeIJacobian(ts,ts->ptime,ts->vec_sol,th->Xdot,shift,J,Jp,PETSC_FALSE);CHKERRQ(ierr); /* get -f_y */
for (nadj=0; nadj<ts->numcost; nadj++) {
ierr = MatMultTranspose(J,VecsDeltaLam[nadj],VecsSensiTemp[nadj]);CHKERRQ(ierr);
ierr = VecAXPY(ts->vecs_sensi[nadj],ts->time_step,VecsSensiTemp[nadj]);CHKERRQ(ierr);
if (ts->vec_costintegral) {
ierr = VecAXPY(ts->vecs_sensi[nadj],-ts->time_step,ts->vecs_drdy[nadj]);CHKERRQ(ierr);
}
}
}
if (ts->vecs_sensip) { /* sensitivities wrt parameters */
ierr = TSAdjointComputeRHSJacobian(ts,ts->ptime,ts->vec_sol,ts->Jacp);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],-ts->time_step*th->Theta,VecsDeltaMu[nadj]);CHKERRQ(ierr);
}
if (th->Theta!=1.) {
ierr = TSAdjointComputeRHSJacobian(ts,th->ptime,th->X0,ts->Jacp);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],-ts->time_step*(1.-th->Theta),VecsDeltaMu[nadj]);CHKERRQ(ierr);
}
}
if (ts->vec_costintegral) {
ierr = TSAdjointComputeDRDPFunction(ts,ts->ptime,ts->vec_sol,ts->vecs_drdp);CHKERRQ(ierr);
for (nadj=0; nadj<ts->numcost; nadj++) {
ierr = VecAXPY(ts->vecs_sensip[nadj],-ts->time_step*th->Theta,ts->vecs_drdp[nadj]);CHKERRQ(ierr);
}
if (th->Theta!=1.) {
ierr = TSAdjointComputeDRDPFunction(ts,th->ptime,th->X0,ts->vecs_drdp);CHKERRQ(ierr);
for (nadj=0; nadj<ts->numcost; nadj++) {
ierr = VecAXPY(ts->vecs_sensip[nadj],-ts->time_step*(1.-th->Theta),ts->vecs_drdp[nadj]);CHKERRQ(ierr);
}
}
}
}
}else { /* one-stage case */
shift = 0.0;
ierr = TSComputeIJacobian(ts,th->stage_time,th->X,th->Xdot,shift,J,Jp,PETSC_FALSE);CHKERRQ(ierr); /* get -f_y */
if (ts->vec_costintegral) {
ierr = TSAdjointComputeDRDYFunction(ts,th->stage_time,th->X,ts->vecs_drdy);CHKERRQ(ierr);
}
for (nadj=0; nadj<ts->numcost; nadj++) {
ierr = MatMultTranspose(J,VecsDeltaLam[nadj],VecsSensiTemp[nadj]);CHKERRQ(ierr);
ierr = VecAXPY(ts->vecs_sensi[nadj],ts->time_step,VecsSensiTemp[nadj]);CHKERRQ(ierr);
if (ts->vec_costintegral) {
ierr = VecAXPY(ts->vecs_sensi[nadj],-ts->time_step,ts->vecs_drdy[nadj]);CHKERRQ(ierr);
}
}
if (ts->vecs_sensip) {
ierr = TSAdjointComputeRHSJacobian(ts,th->stage_time,th->X,ts->Jacp);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],-ts->time_step,VecsDeltaMu[nadj]);CHKERRQ(ierr);
}
if (ts->vec_costintegral) {
ierr = TSAdjointComputeDRDPFunction(ts,th->stage_time,th->X,ts->vecs_drdp);CHKERRQ(ierr);
for (nadj=0; nadj<ts->numcost; nadj++) {
ierr = VecAXPY(ts->vecs_sensip[nadj],-ts->time_step,ts->vecs_drdp[nadj]);CHKERRQ(ierr);
}
}
}
}
th->status = TS_STEP_COMPLETE;
PetscFunctionReturn(0);
}
