static PetscErrorCode TSAdjointStep_RK(TS ts) { TS_RK *rk = (TS_RK*)ts->data; RKTableau tab = rk->tableau; const PetscInt s = tab->s; const PetscReal *A = tab->A,*b = tab->b,*c = tab->c; PetscScalar *w = rk->work; Vec *Y = rk->Y,*VecDeltaLam = rk->VecDeltaLam,*VecDeltaMu = rk->VecDeltaMu,*VecSensiTemp = rk->VecSensiTemp; PetscInt i,j,nadj; PetscReal t = ts->ptime; PetscErrorCode ierr; PetscReal h = ts->time_step; Mat J,Jp; PetscFunctionBegin; rk->status = TS_STEP_INCOMPLETE; for (i=s-1; i>=0; i--) { rk->stage_time = t + h*(1.0-c[i]); for (nadj=0; nadj<ts->numcost; nadj++) { ierr = VecCopy(ts->vecs_sensi[nadj],VecSensiTemp[nadj]);CHKERRQ(ierr); ierr = VecScale(VecSensiTemp[nadj],-h*b[i]);CHKERRQ(ierr); for (j=i+1; j<s; j++) { ierr = VecAXPY(VecSensiTemp[nadj],-h*A[j*s+i],VecDeltaLam[nadj*s+j]);CHKERRQ(ierr); } } /* Stage values of lambda */ ierr = TSGetRHSJacobian(ts,&J,&Jp,NULL,NULL);CHKERRQ(ierr); ierr = TSComputeRHSJacobian(ts,rk->stage_time,Y[i],J,Jp);CHKERRQ(ierr); if (ts->vec_costintegral) { ierr = TSAdjointComputeDRDYFunction(ts,rk->stage_time,Y[i],ts->vecs_drdy);CHKERRQ(ierr); } for (nadj=0; nadj<ts->numcost; nadj++) { ierr = MatMultTranspose(J,VecSensiTemp[nadj],VecDeltaLam[nadj*s+i]);CHKERRQ(ierr); if (ts->vec_costintegral) { ierr = VecAXPY(VecDeltaLam[nadj*s+i],-h*b[i],ts->vecs_drdy[nadj]);CHKERRQ(ierr); } } /* Stage values of mu */ if(ts->vecs_sensip) { ierr = TSAdjointComputeRHSJacobian(ts,rk->stage_time,Y[i],ts->Jacp);CHKERRQ(ierr); if (ts->vec_costintegral) { ierr = TSAdjointComputeDRDPFunction(ts,rk->stage_time,Y[i],ts->vecs_drdp);CHKERRQ(ierr); } for (nadj=0; nadj<ts->numcost; nadj++) { ierr = MatMultTranspose(ts->Jacp,VecSensiTemp[nadj],VecDeltaMu[nadj*s+i]);CHKERRQ(ierr); if (ts->vec_costintegral) { ierr = VecAXPY(VecDeltaMu[nadj*s+i],-h*b[i],ts->vecs_drdp[nadj]);CHKERRQ(ierr); } } } } for (j=0; j<s; j++) w[j] = 1.0; for (nadj=0; nadj<ts->numcost; nadj++) { ierr = VecMAXPY(ts->vecs_sensi[nadj],s,w,&VecDeltaLam[nadj*s]);CHKERRQ(ierr); if(ts->vecs_sensip) { ierr = VecMAXPY(ts->vecs_sensip[nadj],s,w,&VecDeltaMu[nadj*s]);CHKERRQ(ierr); } } rk->status = TS_STEP_COMPLETE; 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); }