Commits

Debojyoti Ghosh  committed 9be3e28

Added a post-stage function for time-integrators

  • Participants
  • Parent commits 9a7421d

Comments (0)

Files changed (11)

File include/petsc-private/tsimpl.h

 
   PetscErrorCode (*prestep)(TS);
   PetscErrorCode (*prestage)(TS,PetscReal);
+  PetscErrorCode (*poststage)(TS,PetscReal,PetscInt,Vec*);
   PetscErrorCode (*poststep)(TS);
 
   /* ---------------------- IMEX support ---------------------------------*/

File include/petscts.h

 
 PETSC_EXTERN PetscErrorCode TSSetPreStep(TS, PetscErrorCode (*)(TS));
 PETSC_EXTERN PetscErrorCode TSSetPreStage(TS, PetscErrorCode (*)(TS,PetscReal));
+PETSC_EXTERN PetscErrorCode TSSetPostStage(TS, PetscErrorCode (*)(TS,PetscReal,PetscInt,Vec*));
 PETSC_EXTERN PetscErrorCode TSSetPostStep(TS, PetscErrorCode (*)(TS));
 PETSC_EXTERN PetscErrorCode TSPreStep(TS);
 PETSC_EXTERN PetscErrorCode TSPreStage(TS,PetscReal);
+PETSC_EXTERN PetscErrorCode TSPostStage(TS,PetscReal,PetscInt,Vec*);
 PETSC_EXTERN PetscErrorCode TSPostStep(TS);
 PETSC_EXTERN PetscErrorCode TSSetRetainStages(TS,PetscBool);
 PETSC_EXTERN PetscErrorCode TSInterpolate(TS,PetscReal,Vec);

File src/ts/impls/arkimex/arkimex.c

     PetscReal h = ts->time_step;
     ierr = TSPreStep(ts);CHKERRQ(ierr);
     for (i=0; i<s; i++) {
+      ark->stage_time = t + h*ct[i];
       if (At[i*s+i] == 0) {           /* This stage is explicit */
         ierr = VecCopy(ts->vec_sol,Y[i]);CHKERRQ(ierr);
         for (j=0; j<i; j++) w[j] = h*At[i*s+j];
         for (j=0; j<i; j++) w[j] = h*A[i*s+j];
         ierr = VecMAXPY(Y[i],i,w,YdotRHS);CHKERRQ(ierr);
       } else {
-        ark->stage_time = t + h*ct[i];
         ark->scoeff     = 1./At[i*s+i];
         ierr            = TSPreStage(ts,ark->stage_time);CHKERRQ(ierr);
         /* Affine part */
         ierr          = TSAdaptCheckStage(adapt,ts,&accept);CHKERRQ(ierr);
         if (!accept) goto reject_step;
       }
+      ierr = TSPostStage(ts,ark->stage_time,i,Y); CHKERRQ(ierr);
       if (ts->equation_type>=TS_EQ_IMPLICIT) {
         if (i==0 && tab->explicit_first_stage) {
           ierr = VecCopy(Ydot0,YdotI[0]);CHKERRQ(ierr);

File src/ts/impls/explicit/euler/euler.c

   ierr = TSPreStage(ts,ts->ptime);CHKERRQ(ierr);
   ierr = TSComputeRHSFunction(ts,ts->ptime,sol,update);CHKERRQ(ierr);
   ierr = VecAXPY(sol,ts->time_step,update);CHKERRQ(ierr);
+  ierr = TSPostStage(ts,ts->ptime,0,&sol);CHKERRQ(ierr);
   ts->ptime += ts->time_step;
   ts->steps++;
   PetscFunctionReturn(0);

File src/ts/impls/explicit/rk/rk.c

     PetscReal h = ts->time_step;
     ierr = TSPreStep(ts);CHKERRQ(ierr);
     for (i=0; i<s; i++) {
+      rk->stage_time = t + h*c[i];
+      ierr = TSPreStage(ts,rk->stage_time); CHKERRQ(ierr);
       ierr = VecCopy(ts->vec_sol,Y[i]);CHKERRQ(ierr);
       for (j=0; j<i; j++) w[j] = h*A[i*s+j];
       ierr = VecMAXPY(Y[i],i,w,YdotRHS);CHKERRQ(ierr);
+      ierr = TSPostStage(ts,rk->stage_time,i,Y); CHKERRQ(ierr);
       ierr = TSGetAdapt(ts,&adapt);CHKERRQ(ierr);
       ierr = TSAdaptCheckStage(adapt,ts,&accept);CHKERRQ(ierr);
       if (!accept) goto reject_step;

File src/ts/impls/implicit/alpha/alpha.c

     /* V1 = (1-1/Gamma)*V0 + 1/(Gamma*dT)*(X1-X0) */
     ierr = VecWAXPY(th->V1,-1,th->X0,th->X1);CHKERRQ(ierr);
     ierr = VecAXPBY(th->V1,1-1/th->Gamma,1/(th->Gamma*ts->time_step),th->V0);CHKERRQ(ierr);
+    ierr = TSPostStage(ts,th->stage_time,0,&(th->V1));CHKERRQ(ierr);
     /* nonlinear solve convergence */
     ierr = SNESGetConvergedReason(ts->snes,&snesreason);CHKERRQ(ierr);
     if (snesreason < 0 && !th->adapt) break;

File src/ts/impls/implicit/sundials/sundials.c

 
   ierr = TSPreStep(ts);CHKERRQ(ierr);
 
-  /* We would like to call TSPreStep() when starting each step (including rejections) and TSPreStage() before each
-   * stage solve, but CVode does not appear to support this. */
+  /* We would like to call TSPreStep() when starting each step (including rejections), TSPreStage(),
+   * and TSPostStage() before each stage solve, but CVode does not appear to support this. */
   if (cvode->monitorstep) flag = CVode(mem,tout,cvode->y,&t,CV_ONE_STEP);
   else flag = CVode(mem,tout,cvode->y,&t,CV_NORMAL);
 

File src/ts/impls/implicit/theta/theta.c

     ierr = SNESGetIterationNumber(ts->snes,&its);CHKERRQ(ierr);
     ierr = SNESGetLinearSolveIterations(ts->snes,&lits);CHKERRQ(ierr);
     ierr = SNESGetConvergedReason(ts->snes,&snesreason);CHKERRQ(ierr);
+    ierr = TSPostStage(ts,th->stage_time,0,&(th->X));CHKERRQ(ierr);
     ts->snes_its += its; ts->ksp_its += lits;
     ierr = TSGetAdapt(ts,&adapt);CHKERRQ(ierr);
     ierr = TSAdaptCheckStage(adapt,ts,&accept);CHKERRQ(ierr);

File src/ts/impls/pseudo/posindep.c

     ierr = SNESGetConvergedReason(ts->snes,&snesreason);CHKERRQ(ierr);
     ierr = SNESGetLinearSolveIterations(ts->snes,&lits);CHKERRQ(ierr);
     ierr = SNESGetIterationNumber(ts->snes,&its);CHKERRQ(ierr);
+    ierr = TSPostStage(ts,ts->ptime+ts->time_step,0,&(pseudo->update));CHKERRQ(ierr);
     ts->snes_its += its; ts->ksp_its += lits;
     ierr = PetscInfo3(ts,"step=%D, nonlinear solve iterations=%D, linear solve iterations=%D\n",ts->steps,its,lits);CHKERRQ(ierr);
     pseudo->fnorm = -1;         /* The current norm is no longer valid, monitor must recompute it. */

File src/ts/impls/rosw/rosw.c

         ierr = VecScale(Y[i],h);
         ts->ksp_its += 1;
       }
+      ierr = TSPostStage(ts,ros->stage_time,i,Y);CHKERRQ(ierr);
     }
     ierr = TSEvaluateStep(ts,tab->order,ts->vec_sol,NULL);CHKERRQ(ierr);
     ros->status = TS_STEP_PENDING;

File src/ts/interface/ts.c

    Level: intermediate
 
 .keywords: TS, timestep, get, iteration, number
-.seealso: TSGetTime(), TSGetTimeStep(), TSSetPreStep(), TSSetPreStage(), TSSetPostStep()
+.seealso: TSGetTime(), TSGetTimeStep(), TSSetPreStep(), TSSetPreStage(), TSSetPostStage(), TSSetPostStep()
 @*/
 PetscErrorCode  TSGetTimeStepNumber(TS ts,PetscInt *iter)
 {
   size of the step being attempted can be obtained using TSGetTimeStep().
 
 .keywords: TS, timestep
-.seealso: TSSetPreStage(), TSSetPostStep(), TSStep()
+.seealso: TSSetPreStage(), TSSetPostStage(), TSSetPostStep(), TSStep()
 @*/
 PetscErrorCode  TSSetPreStep(TS ts, PetscErrorCode (*func)(TS))
 {
   Level: developer
 
 .keywords: TS, timestep
-.seealso: TSSetPreStep(), TSPreStage(), TSPostStep()
+.seealso: TSSetPreStep(), TSPreStage(), TSPostStage(), TSPostStep()
 @*/
 PetscErrorCode  TSPreStep(TS ts)
 {
   attempted can be obtained using TSGetTimeStep(). The time at the start of the step is available via TSGetTime().
 
 .keywords: TS, timestep
-.seealso: TSSetPreStep(), TSSetPostStep(), TSGetApplicationContext()
+.seealso: TSSetPostStage(), TSSetPreStep(), TSSetPostStep(), TSGetApplicationContext()
 @*/
 PetscErrorCode  TSSetPreStage(TS ts, PetscErrorCode (*func)(TS,PetscReal))
 {
 }
 
 #undef __FUNCT__
+#define __FUNCT__ "TSSetPostStage"
+/*@C
+  TSSetPostStage - Sets the general-purpose function
+  called once at the end of each stage.
+
+  Logically Collective on TS
+
+  Input Parameters:
++ ts   - The TS context obtained from TSCreate()
+- func - The function
+
+  Calling sequence of func:
+. PetscErrorCode func(TS ts, PetscReal stagetime, PetscInt stageindex, Vec* Y);
+
+  Level: intermediate
+
+  Note:
+  There may be several stages per time step. If the solve for a given stage fails, the step may be rejected and retried.
+  The time step number being computed can be queried using TSGetTimeStepNumber() and the total size of the step being
+  attempted can be obtained using TSGetTimeStep(). The time at the start of the step is available via TSGetTime().
+
+.keywords: TS, timestep
+.seealso: TSSetPreStage(), TSSetPreStep(), TSSetPostStep(), TSGetApplicationContext()
+@*/
+PetscErrorCode  TSSetPostStage(TS ts, PetscErrorCode (*func)(TS,PetscReal,PetscInt,Vec*))
+{
+  PetscFunctionBegin;
+  PetscValidHeaderSpecific(ts, TS_CLASSID,1);
+  ts->poststage = func;
+  PetscFunctionReturn(0);
+}
+
+#undef __FUNCT__
 #define __FUNCT__ "TSPreStage"
 /*@
   TSPreStage - Runs the user-defined pre-stage function set using TSSetPreStage()
   Collective on TS
 
   Input Parameters:
-. ts   - The TS context obtained from TSCreate()
+. ts          - The TS context obtained from TSCreate()
+  stagetime   - The absolute time of the current stage
 
   Notes:
   TSPreStage() is typically used within time stepping implementations,
   Level: developer
 
 .keywords: TS, timestep
-.seealso: TSSetPreStep(), TSPreStep(), TSPostStep()
+.seealso: TSPostStage(), TSSetPreStep(), TSPreStep(), TSPostStep()
 @*/
 PetscErrorCode  TSPreStage(TS ts, PetscReal stagetime)
 {
 }
 
 #undef __FUNCT__
+#define __FUNCT__ "TSPostStage"
+/*@
+  TSPostStage - Runs the user-defined post-stage function set using TSSetPostStage()
+
+  Collective on TS
+
+  Input Parameters:
+. ts          - The TS context obtained from TSCreate()
+  stagetime   - The absolute time of the current stage
+  stageindex  - Stage number
+  Y           - Array of vectors (of size = total number
+                of stages) with the stage solutions
+
+  Notes:
+  TSPostStage() is typically used within time stepping implementations,
+  most users would not generally call this routine themselves.
+
+  Level: developer
+
+.keywords: TS, timestep
+.seealso: TSPreStage(), TSSetPreStep(), TSPreStep(), TSPostStep()
+@*/
+PetscErrorCode  TSPostStage(TS ts, PetscReal stagetime, PetscInt stageindex, Vec *Y)
+{
+  PetscErrorCode ierr;
+
+  PetscFunctionBegin;
+  PetscValidHeaderSpecific(ts,TS_CLASSID,1);
+  if (ts->prestage) {
+    PetscStackCallStandard((*ts->poststage),(ts,stagetime,stageindex,Y));
+  }
+  PetscFunctionReturn(0);
+}
+
+#undef __FUNCT__
 #define __FUNCT__ "TSSetPostStep"
 /*@C
   TSSetPostStep - Sets the general-purpose function
 
 .keywords: TS, timestep, solve
 
-.seealso: TSCreate(), TSSetUp(), TSDestroy(), TSSolve(), TSSetPreStep(), TSSetPreStage(), TSInterpolate()
+.seealso: TSCreate(), TSSetUp(), TSDestroy(), TSSolve(), TSSetPreStep(), TSSetPreStage(), TSSetPostStage(), TSInterpolate()
 @*/
 PetscErrorCode  TSStep(TS ts)
 {
 
    Note:
    When called during time step evaluation (e.g. during residual evaluation or via hooks set using TSSetPreStep(),
-   TSSetPreStage(), or TSSetPostStep()), the time is the time at the start of the step being evaluated.
+   TSSetPreStage(), TSSetPostStage(), or TSSetPostStep()), the time is the time at the start of the step being evaluated.
 
 .seealso: TSSetInitialTimeStep(), TSGetTimeStep()