static PetscErrorCode TSStep_Alpha(TS ts)
{
TS_Alpha *th = (TS_Alpha*)ts->data;
PetscInt rejections = 0;
PetscBool stageok,accept = PETSC_TRUE;
PetscReal next_time_step = ts->time_step;
PetscErrorCode ierr;
PetscFunctionBegin;
ierr = PetscCitationsRegister(citation,&cited);CHKERRQ(ierr);
if (!ts->steprollback) {
if (th->adapt) { ierr = VecCopy(th->X0,th->vec_sol_prev);CHKERRQ(ierr); }
if (th->adapt) { ierr = VecCopy(th->V0,th->vec_dot_prev);CHKERRQ(ierr); }
ierr = VecCopy(ts->vec_sol,th->X0);CHKERRQ(ierr);
ierr = VecCopy(ts->vec_dot,th->V0);CHKERRQ(ierr);
ierr = VecCopy(th->A1,th->A0);CHKERRQ(ierr);
}
th->status = TS_STEP_INCOMPLETE;
while (!ts->reason && th->status != TS_STEP_COMPLETE) {
if (ts->steprestart) {
ierr = TSAlpha_Restart(ts,&stageok);CHKERRQ(ierr);
if (!stageok) goto reject_step;
}
ierr = TSAlpha_StageTime(ts);CHKERRQ(ierr);
ierr = VecCopy(th->X0,th->X1);CHKERRQ(ierr);
ierr = TSPreStage(ts,th->stage_time);CHKERRQ(ierr);
ierr = TS_SNESSolve(ts,NULL,th->X1);CHKERRQ(ierr);
ierr = TSPostStage(ts,th->stage_time,0,&th->Xa);CHKERRQ(ierr);
ierr = TSAdaptCheckStage(ts->adapt,ts,th->stage_time,th->Xa,&stageok);CHKERRQ(ierr);
if (!stageok) goto reject_step;
th->status = TS_STEP_PENDING;
ierr = VecCopy(th->X1,ts->vec_sol);CHKERRQ(ierr);
ierr = VecCopy(th->V1,ts->vec_dot);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);
ierr = VecCopy(th->V0,ts->vec_dot);CHKERRQ(ierr);
ts->time_step = next_time_step;
goto reject_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 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->adapt) { 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) {
PetscReal 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/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 = TS_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,-shift,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->costintegralfwd) { /* Save the info for the later use in cost integral evaluation */
th->ptime = ts->ptime;
th->time_step = 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 TSAlpha_Restart(TS ts,PetscBool *initok)
{
TS_Alpha *th = (TS_Alpha*)ts->data;
PetscReal time_step;
PetscReal alpha_m,alpha_f,gamma,beta;
Vec X0 = ts->vec_sol, X1, X2 = th->X1;
Vec V0 = ts->vec_dot, V1, V2 = th->V1;
PetscBool stageok;
PetscErrorCode ierr;
PetscFunctionBegin;
ierr = VecDuplicate(X0,&X1);CHKERRQ(ierr);
ierr = VecDuplicate(V0,&V1);CHKERRQ(ierr);
/* Setup backward Euler with halved time step */
ierr = TSAlpha2GetParams(ts,&alpha_m,&alpha_f,&gamma,&beta);CHKERRQ(ierr);
ierr = TSAlpha2SetParams(ts,1,1,1,0.5);CHKERRQ(ierr);
ierr = TSGetTimeStep(ts,&time_step);CHKERRQ(ierr);
ts->time_step = time_step/2;
ierr = TSAlpha_StageTime(ts);CHKERRQ(ierr);
th->stage_time = ts->ptime;
ierr = VecZeroEntries(th->A0);CHKERRQ(ierr);
/* First BE step, (t0,X0,V0) -> (t1,X1,V1) */
th->stage_time += ts->time_step;
ierr = VecCopy(X0,th->X0);CHKERRQ(ierr);
ierr = VecCopy(V0,th->V0);CHKERRQ(ierr);
ierr = TSPreStage(ts,th->stage_time);CHKERRQ(ierr);
ierr = VecCopy(th->X0,X1);CHKERRQ(ierr);
ierr = TS_SNESSolve(ts,NULL,X1);CHKERRQ(ierr);
ierr = VecCopy(th->V1,V1);CHKERRQ(ierr);
ierr = TSPostStage(ts,th->stage_time,0,&X1);CHKERRQ(ierr);
ierr = TSAdaptCheckStage(ts->adapt,ts,th->stage_time,X1,&stageok);CHKERRQ(ierr);
if (!stageok) goto finally;
/* Second BE step, (t1,X1,V1) -> (t2,X2,V2) */
th->stage_time += ts->time_step;
ierr = VecCopy(X1,th->X0);CHKERRQ(ierr);
ierr = VecCopy(V1,th->V0);CHKERRQ(ierr);
ierr = TSPreStage(ts,th->stage_time);CHKERRQ(ierr);
ierr = VecCopy(th->X0,X2);CHKERRQ(ierr);
ierr = TS_SNESSolve(ts,NULL,X2);CHKERRQ(ierr);
ierr = VecCopy(th->V1,V2);CHKERRQ(ierr);
ierr = TSPostStage(ts,th->stage_time,0,&X2);CHKERRQ(ierr);
ierr = TSAdaptCheckStage(ts->adapt,ts,th->stage_time,X1,&stageok);CHKERRQ(ierr);
if (!stageok) goto finally;
/* Compute A0 ~ dV/dt at t0 with backward differences */
ierr = VecZeroEntries(th->A0);CHKERRQ(ierr);
ierr = VecAXPY(th->A0,-3/ts->time_step,V0);CHKERRQ(ierr);
ierr = VecAXPY(th->A0,+4/ts->time_step,V1);CHKERRQ(ierr);
ierr = VecAXPY(th->A0,-1/ts->time_step,V2);CHKERRQ(ierr);
/* Rough, lower-order estimate LTE of the initial step */
if (th->adapt) {
ierr = VecZeroEntries(th->vec_lte_work[0]);CHKERRQ(ierr);
ierr = VecAXPY(th->vec_lte_work[0],+2,X2);CHKERRQ(ierr);
ierr = VecAXPY(th->vec_lte_work[0],-4,X1);CHKERRQ(ierr);
ierr = VecAXPY(th->vec_lte_work[0],+2,X0);CHKERRQ(ierr);
}
if (th->adapt) {
ierr = VecZeroEntries(th->vec_lte_work[1]);CHKERRQ(ierr);
ierr = VecAXPY(th->vec_lte_work[1],+2,V2);CHKERRQ(ierr);
ierr = VecAXPY(th->vec_lte_work[1],-4,V1);CHKERRQ(ierr);
ierr = VecAXPY(th->vec_lte_work[1],+2,V0);CHKERRQ(ierr);
}
finally:
/* Revert TSAlpha to the initial state (t0,X0,V0) */
if (initok) *initok = stageok;
ierr = TSSetTimeStep(ts,time_step);CHKERRQ(ierr);
ierr = TSAlpha2SetParams(ts,alpha_m,alpha_f,gamma,beta);CHKERRQ(ierr);
ierr = VecCopy(ts->vec_sol,th->X0);CHKERRQ(ierr);
ierr = VecCopy(ts->vec_dot,th->V0);CHKERRQ(ierr);
ierr = VecDestroy(&X1);CHKERRQ(ierr);
ierr = VecDestroy(&V1);CHKERRQ(ierr);
PetscFunctionReturn(0);
}
