xref: /petsc/src/ts/trajectory/impls/memory/trajmemory.c (revision 42f060c66c5139b47b918cdcc5dbde6c2c2c2c65)
1 #define TJ_VERBOSE
2 #include <petsc/private/tsimpl.h>        /*I "petscts.h"  I*/
3 #include <petscsys.h>
4 
5 extern int wrap_revolve(int* check,int* capo,int* fine,int *snaps_in,int* info,int* rank);
6 
7 typedef struct _StackElement {
8   PetscInt  stepnum;
9   Vec       X;
10   Vec       *Y;
11   PetscReal time;
12   PetscReal timeprev;
13 } *StackElement;
14 
15 typedef struct _RevolveCTX {
16   PetscBool    reverseonestep;
17   PetscInt     snaps_in;
18   PetscInt     stepsleft;
19   PetscInt     check;
20   PetscInt     oldcapo;
21   PetscInt     capo;
22   PetscInt     fine;
23   PetscInt     info;
24 } RevolveCTX;
25 
26 typedef struct _Stack {
27   PetscBool    userevolve;
28   RevolveCTX   *rctx;
29   PetscInt     top;         /* top of the stack */
30   PetscInt     max_cps;     /* maximum stack size */
31   PetscInt     numY;
32   PetscInt     stride;
33   PetscInt     total_steps; /* total number of steps */
34   MPI_Comm     comm;
35   StackElement *stack;      /* container */
36 } Stack;
37 
38 static PetscErrorCode StackCreate(MPI_Comm,Stack *,PetscInt,PetscInt);
39 static PetscErrorCode StackDestroy(Stack*);
40 static PetscErrorCode StackPush(Stack*,StackElement);
41 static PetscErrorCode StackPop(Stack*,StackElement*);
42 static PetscErrorCode StackTop(Stack*,StackElement*);
43 static PetscErrorCode StackDumpAll(TS,Stack*,PetscInt);
44 static PetscErrorCode StackLoadAll(TS,Stack*,PetscInt);
45 
46 #ifdef TJ_VERBOSE
47 static void printwhattodo(PetscInt whattodo,RevolveCTX *rctx,PetscInt shift)
48 {
49   switch(whattodo) {
50     case 1:
51       PetscPrintf(PETSC_COMM_WORLD,"Advance from %D to %D.\n",rctx->oldcapo+shift,rctx->capo+shift);
52       break;
53     case 2:
54       PetscPrintf(PETSC_COMM_WORLD,"Store in checkpoint number %D\n",rctx->check);
55       break;
56     case 3:
57       PetscPrintf(PETSC_COMM_WORLD,"First turn: Initialize adjoints and reverse first step.\n");
58       break;
59     case 4:
60       PetscPrintf(PETSC_COMM_WORLD,"Forward and reverse one step.\n");
61       break;
62     case 5:
63       PetscPrintf(PETSC_COMM_WORLD,"Restore in checkpoint number %D\n",rctx->check);
64       break;
65     case -1:
66       PetscPrintf(PETSC_COMM_WORLD,"Error!");
67       break;
68   }
69 }
70 #endif
71 
72 #undef __FUNCT__
73 #define __FUNCT__ "StackCreate"
74 static PetscErrorCode StackCreate(MPI_Comm comm,Stack *s,PetscInt size,PetscInt ny)
75 {
76   PetscErrorCode ierr;
77 
78   PetscFunctionBegin;
79   s->top     = -1;
80   s->max_cps = size;
81   s->comm    = comm;
82   s->numY    = ny;
83 
84   ierr = PetscMalloc1(s->max_cps*sizeof(StackElement),&s->stack);CHKERRQ(ierr);
85   PetscFunctionReturn(0);
86 }
87 
88 #undef __FUNCT__
89 #define __FUNCT__ "StackDestroy"
90 static PetscErrorCode StackDestroy(Stack *s)
91 {
92   PetscInt       i;
93   PetscErrorCode ierr;
94 
95   PetscFunctionBegin;
96   if (s->top>-1) {
97     for (i=0;i<=s->top;i++) {
98       ierr = VecDestroy(&s->stack[i]->X);CHKERRQ(ierr);
99       ierr = VecDestroyVecs(s->numY,&s->stack[i]->Y);CHKERRQ(ierr);
100       ierr = PetscFree(s->stack[i]);CHKERRQ(ierr);
101     }
102   }
103   ierr = PetscFree(s->stack);CHKERRQ(ierr);
104   if (s->userevolve) {
105     ierr = PetscFree(s->rctx);CHKERRQ(ierr);
106   }
107   ierr = PetscFree(s);CHKERRQ(ierr);
108   PetscFunctionReturn(0);
109 }
110 
111 #undef __FUNCT__
112 #define __FUNCT__ "StackPush"
113 static PetscErrorCode StackPush(Stack *s,StackElement e)
114 {
115   PetscFunctionBegin;
116   if (s->userevolve && s->top+1 >= s->max_cps) SETERRQ1(s->comm,PETSC_ERR_MEMC,"Maximum stack size (%D) exceeded",s->max_cps);
117   if (!s->userevolve && s->top+1 >= s->stride-1) SETERRQ1(s->comm,PETSC_ERR_MEMC,"Maximum stack size (%D) exceeded",s->max_cps);
118   s->stack[++s->top] = e;
119   PetscFunctionReturn(0);
120 }
121 
122 #undef __FUNCT__
123 #define __FUNCT__ "StackPop"
124 static PetscErrorCode StackPop(Stack *s,StackElement *e)
125 {
126   PetscFunctionBegin;
127   if (s->top == -1) SETERRQ(s->comm,PETSC_ERR_MEMC,"Emptry stack");
128   *e = s->stack[s->top--];
129   PetscFunctionReturn(0);
130 }
131 
132 #undef __FUNCT__
133 #define __FUNCT__ "StackTop"
134 static PetscErrorCode StackTop(Stack *s,StackElement *e)
135 {
136   PetscFunctionBegin;
137   *e = s->stack[s->top];
138   PetscFunctionReturn(0);
139 }
140 
141 #undef __FUNCT__
142 #define __FUNCT__ "OutputBIN"
143 static PetscErrorCode OutputBIN(const char *filename,PetscViewer *viewer)
144 {
145   PetscErrorCode ierr;
146 
147   PetscFunctionBegin;
148   ierr = PetscViewerCreate(PETSC_COMM_WORLD,viewer);CHKERRQ(ierr);
149   ierr = PetscViewerSetType(*viewer,PETSCVIEWERBINARY);CHKERRQ(ierr);
150   ierr = PetscViewerFileSetMode(*viewer,FILE_MODE_WRITE);CHKERRQ(ierr);
151   ierr = PetscViewerFileSetName(*viewer,filename);CHKERRQ(ierr);
152   PetscFunctionReturn(0);
153 }
154 
155 #undef __FUNCT__
156 #define __FUNCT__ "StackDumpAll"
157 static PetscErrorCode StackDumpAll(TS ts,Stack *s,PetscInt id)
158 {
159   PetscInt       i,j;
160   StackElement   e;
161   PetscViewer    viewer;
162   char           filename[PETSC_MAX_PATH_LEN];
163   Vec            *Y;
164   PetscErrorCode ierr;
165 
166   PetscFunctionBegin;
167   if (id == 1) {
168 #if defined(PETSC_HAVE_POPEN)
169     PetscMPIInt rank;
170     ierr = MPI_Comm_rank(s->comm,&rank);CHKERRQ(ierr);
171     if (!rank) {
172       char command[PETSC_MAX_PATH_LEN];
173       FILE *fd;
174       int  err;
175 
176       ierr = PetscMemzero(command,sizeof(command));CHKERRQ(ierr);
177       ierr = PetscSNPrintf(command,PETSC_MAX_PATH_LEN,"rm -fr %s","SA-data");CHKERRQ(ierr);
178       ierr = PetscPOpen(PETSC_COMM_SELF,NULL,command,"r",&fd);CHKERRQ(ierr);
179       ierr = PetscPClose(PETSC_COMM_SELF,fd,&err);CHKERRQ(ierr);
180       ierr = PetscSNPrintf(command,PETSC_MAX_PATH_LEN,"mkdir %s","SA-data");CHKERRQ(ierr);
181       ierr = PetscPOpen(PETSC_COMM_SELF,NULL,command,"r",&fd);CHKERRQ(ierr);
182       ierr = PetscPClose(PETSC_COMM_SELF,fd,&err);CHKERRQ(ierr);
183     }
184 #endif
185   }
186   ierr = PetscSNPrintf(filename,sizeof(filename),"SA-data/SA-STACK%06d.bin",id);CHKERRQ(ierr);
187   ierr = OutputBIN(filename,&viewer);CHKERRQ(ierr);
188   for (i=0;i<s->stride-1;i++) {
189     e = s->stack[i];
190     ierr = PetscViewerBinaryWrite(viewer,&e->stepnum,1,PETSC_INT,PETSC_FALSE);CHKERRQ(ierr);
191     ierr = VecView(e->X,viewer);CHKERRQ(ierr);
192     for (j=0;j<s->numY;j++) {
193       ierr = VecView(e->Y[j],viewer);CHKERRQ(ierr);
194     }
195     ierr = PetscViewerBinaryWrite(viewer,&e->time,1,PETSC_REAL,PETSC_FALSE);CHKERRQ(ierr);
196     ierr = PetscViewerBinaryWrite(viewer,&e->timeprev,1,PETSC_REAL,PETSC_FALSE);CHKERRQ(ierr);
197   }
198   /* dump the state inside TS from the current step */
199   ierr = PetscViewerBinaryWrite(viewer,&ts->total_steps,1,PETSC_INT,PETSC_FALSE);CHKERRQ(ierr);
200   ierr = VecView(ts->vec_sol,viewer);CHKERRQ(ierr);
201   ierr = TSGetStages(ts,&s->numY,&Y);CHKERRQ(ierr);
202   for (j=0;j<s->numY;j++) {
203     ierr = VecView(Y[j],viewer);CHKERRQ(ierr);
204   }
205   ierr = PetscViewerBinaryWrite(viewer,&ts->ptime,1,PETSC_REAL,PETSC_FALSE);CHKERRQ(ierr);
206   ierr = PetscViewerBinaryWrite(viewer,&ts->ptime_prev,1,PETSC_REAL,PETSC_FALSE);CHKERRQ(ierr);
207   for (i=0;i<s->stride-1;i++) {
208     ierr = StackPop(s,&e);CHKERRQ(ierr);
209     ierr = VecDestroy(&e->X);CHKERRQ(ierr);
210     ierr = VecDestroyVecs(s->numY,&e->Y);CHKERRQ(ierr);
211     ierr = PetscFree(e);CHKERRQ(ierr);
212   }
213   ierr = PetscViewerDestroy(&viewer);CHKERRQ(ierr);
214   PetscFunctionReturn(0);
215 }
216 
217 #undef __FUNCT__
218 #define __FUNCT__ "StackLoadAll"
219 static PetscErrorCode StackLoadAll(TS ts,Stack *s,PetscInt id)
220 {
221   Vec            *Y;
222   PetscInt       i,j;
223   StackElement   e;
224   PetscViewer    viewer;
225   char           filename[PETSC_MAX_PATH_LEN];
226   PetscErrorCode ierr;
227 
228   PetscFunctionBegin;
229   ierr = PetscSNPrintf(filename,sizeof filename,"SA-data/SA-STACK%06d.bin",id);CHKERRQ(ierr);
230   ierr = PetscViewerBinaryOpen(PETSC_COMM_WORLD,filename,FILE_MODE_READ,&viewer);CHKERRQ(ierr);
231   for (i=0;i<s->stride-1;i++) {
232     ierr = PetscCalloc1(1,&e);
233     ierr = TSGetStages(ts,&s->numY,&Y);CHKERRQ(ierr);
234     ierr = VecDuplicate(Y[0],&e->X);CHKERRQ(ierr);
235     ierr = VecDuplicateVecs(Y[0],s->numY,&e->Y);CHKERRQ(ierr);
236     ierr = StackPush(s,e);CHKERRQ(ierr);
237   }
238   for (i=0;i<s->stride-1;i++) {
239     e = s->stack[i];
240     ierr = PetscViewerBinaryRead(viewer,&e->stepnum,1,NULL,PETSC_INT);CHKERRQ(ierr);
241     ierr = VecLoad(e->X,viewer);CHKERRQ(ierr);
242     for (j=0;j<s->numY;j++) {
243       ierr = VecLoad(e->Y[j],viewer);CHKERRQ(ierr);
244     }
245     ierr = PetscViewerBinaryRead(viewer,&e->time,1,NULL,PETSC_REAL);CHKERRQ(ierr);
246     ierr = PetscViewerBinaryRead(viewer,&e->timeprev,1,NULL,PETSC_REAL);CHKERRQ(ierr);
247   }
248   /* load the additioinal state into TS */
249   ierr = PetscViewerBinaryRead(viewer,&ts->total_steps,1,NULL,PETSC_INT);CHKERRQ(ierr);
250   ierr = VecLoad(ts->vec_sol,viewer);CHKERRQ(ierr);
251   ierr = TSGetStages(ts,&s->numY,&Y);CHKERRQ(ierr);
252   for (j=0;j<s->numY;j++) {
253     ierr = VecLoad(Y[j],viewer);CHKERRQ(ierr);
254   }
255   ierr = PetscViewerBinaryRead(viewer,&ts->ptime,1,NULL,PETSC_REAL);CHKERRQ(ierr);
256   ierr = PetscViewerBinaryRead(viewer,&ts->ptime_prev,1,NULL,PETSC_REAL);CHKERRQ(ierr);
257   ierr = TSSetTimeStep(ts,ts->ptime-ts->ptime_prev);CHKERRQ(ierr);
258   ierr = PetscViewerDestroy(&viewer);CHKERRQ(ierr);
259   PetscFunctionReturn(0);
260 }
261 
262 #undef __FUNCT__
263 #define __FUNCT__ "TSTrajectorySetStride_Memory"
264 PetscErrorCode TSTrajectorySetStride_Memory(TSTrajectory tj,TS ts,PetscInt stride)
265 {
266   Stack    *s = (Stack*)tj->data;
267 
268   PetscFunctionBegin;
269   s->stride = stride;
270   PetscFunctionReturn(0);
271 }
272 
273 #undef __FUNCT__
274 #define __FUNCT__ "TSTrajectorySetMaxCheckpoints_Memory"
275 PetscErrorCode TSTrajectorySetMaxCheckpoints_Memory(TSTrajectory tj,TS ts,PetscInt max_cps)
276 {
277   Stack    *s = (Stack*)tj->data;
278 
279   PetscFunctionBegin;
280   s->max_cps = max_cps;
281   PetscFunctionReturn(0);
282 }
283 
284 #undef __FUNCT__
285 #define __FUNCT__ "TSTrajectorySetFromOptions_Memory"
286 PetscErrorCode TSTrajectorySetFromOptions_Memory(PetscOptions *PetscOptionsObject,TSTrajectory tj)
287 {
288   Stack     *s = (Stack*)tj->data;
289   PetscErrorCode ierr;
290 
291   PetscFunctionBegin;
292   ierr = PetscOptionsHead(PetscOptionsObject,"Memory based TS trajectory options");CHKERRQ(ierr);
293   {
294     ierr = PetscOptionsInt("-tstrajectory_max_cps","Maximum number of checkpoints","TSTrajectorySetMaxCheckpoints_Memory",s->max_cps,&s->max_cps,NULL);CHKERRQ(ierr);
295     ierr = PetscOptionsInt("-tstrajectory_stride","Stride to save checkpoints to file","TSTrajectorySetStride_Memory",s->stride,&s->stride,NULL);CHKERRQ(ierr);
296   }
297   ierr = PetscOptionsTail();CHKERRQ(ierr);
298   PetscFunctionReturn(0);
299 }
300 
301 #undef __FUNCT__
302 #define __FUNCT__ "TSTrajectorySetUp_Memory"
303 PetscErrorCode TSTrajectorySetUp_Memory(TSTrajectory tj,TS ts)
304 {
305   Stack          *s = (Stack*)tj->data;
306   RevolveCTX     *rctx;
307   PetscInt       numY;
308   PetscErrorCode ierr;
309 
310   PetscFunctionBegin;
311   ierr = TSGetStages(ts,&numY,NULL);CHKERRQ(ierr);
312   s->total_steps = PetscMin(ts->max_steps,(PetscInt)(ceil(ts->max_time/ts->time_step)));
313 
314   if ((s->stride>1 && s->max_cps>1 && s->max_cps<s->stride-1)||(s->stride<=1 && s->max_cps>1 && s->max_cps<s->total_steps-1)) {
315     s->userevolve  = PETSC_TRUE;
316     ierr = PetscCalloc1(1,&rctx);CHKERRQ(ierr);
317     rctx->snaps_in       = s->max_cps; /* for theta methods snaps_in=2*max_cps */
318     rctx->reverseonestep = PETSC_FALSE;
319     rctx->check          = -1;
320     rctx->oldcapo        = 0;
321     rctx->capo           = 0;
322     rctx->info           = 2;
323     s->rctx = rctx;
324     if (s->stride>1) rctx->fine = s->stride;
325     else rctx->fine = s->total_steps;
326     ierr = StackCreate(PetscObjectComm((PetscObject)ts),s,s->max_cps,numY);CHKERRQ(ierr);
327   } else {
328     s->userevolve = PETSC_FALSE;
329     if (s->stride>1) {
330       ierr = StackCreate(PetscObjectComm((PetscObject)ts),s,s->stride-1,numY);CHKERRQ(ierr);
331     } else {
332       ierr = StackCreate(PetscObjectComm((PetscObject)ts),s,ts->max_steps-1,numY);CHKERRQ(ierr);
333     }
334   }
335   PetscFunctionReturn(0);
336 }
337 
338 #undef __FUNCT__
339 #define __FUNCT__ "TSTrajectorySet_Memory"
340 PetscErrorCode TSTrajectorySet_Memory(TSTrajectory tj,TS ts,PetscInt stepnum,PetscReal time,Vec X)
341 {
342   PetscInt       i;
343   Vec            *Y;
344   PetscReal      timeprev;
345   StackElement   e;
346   Stack          *s = (Stack*)tj->data;
347   RevolveCTX     *rctx;
348   PetscInt       whattodo,rank,localstepnum,id,shift;
349   PetscErrorCode ierr;
350 
351   PetscFunctionBegin;
352   if (stepnum<s->top) SETERRQ(s->comm,PETSC_ERR_MEMC,"Illegal modification of a non-top stack element");
353 
354   if (s->stride>1) { /* multilevel mode */
355     localstepnum = stepnum%s->stride;
356     if (stepnum!=0 && stepnum!=s->total_steps && localstepnum==0) {
357       id = stepnum/s->stride;
358       ierr = StackDumpAll(ts,s,id);CHKERRQ(ierr);
359       s->top = -1; /* reset top */
360     }
361     if (s->userevolve) {
362       rctx = s->rctx;
363       rctx->check = -1;
364       rctx->capo  = 0;
365       rctx->fine  = s->stride;
366     }
367   } else { /* in-memory mode */
368     localstepnum = stepnum;
369   }
370 
371   if (s->userevolve) {
372     rctx = s->rctx;
373     if (rctx->reverseonestep) { /* intermediate information is ready inside TS */
374       rctx->reverseonestep = PETSC_FALSE;
375       PetscFunctionReturn(0);
376     }
377     if (rctx->stepsleft!=0) { /* advance the solution without checkpointing anything as Revolve requires */
378       rctx->stepsleft--;
379       PetscFunctionReturn(0);
380     }
381 
382     /* let Revolve determine what to do next */
383     shift         = stepnum-localstepnum;
384     rctx->capo    = localstepnum;
385     rctx->oldcapo = rctx->capo;
386     whattodo = wrap_revolve(&rctx->check,&rctx->capo,&rctx->fine,&rctx->snaps_in,&rctx->info,&rank);
387 #ifdef TJ_VERBOSE
388     printwhattodo(whattodo,rctx,shift);
389 #endif
390     if (whattodo==-1) SETERRQ(s->comm,PETSC_ERR_MEMC,"Error in the controller");
391     if (whattodo==1) { /* advance some time steps */
392       rctx->stepsleft = rctx->capo-rctx->oldcapo-1;
393       PetscFunctionReturn(0);
394     }
395     if (whattodo==3 || whattodo==4) { /* ready for a reverse step */
396       rctx->reverseonestep = PETSC_TRUE;
397       PetscFunctionReturn(0);
398     }
399     if (whattodo==5) { /* restore a checkpoint and ask Revolve what to do next */
400       rctx->oldcapo = rctx->capo;
401       whattodo = wrap_revolve(&rctx->check,&rctx->capo,&rctx->fine,&rctx->snaps_in,&rctx->info,&rank); /* must return 1*/
402 #ifdef TJ_VERBOSE
403       printwhattodo(whattodo,rctx,shift);
404 #endif
405       rctx->stepsleft = rctx->capo-rctx->oldcapo;
406       PetscFunctionReturn(0);
407     }
408     if (whattodo==2) { /* store a checkpoint and ask Revolve how many time steps to advance next */
409       rctx->oldcapo = rctx->capo;
410       whattodo = wrap_revolve(&rctx->check,&rctx->capo,&rctx->fine,&rctx->snaps_in,&rctx->info,&rank); /* must return 1*/
411 #ifdef TJ_VERBOSE
412       printwhattodo(whattodo,rctx,shift);
413 #endif
414       rctx->stepsleft = rctx->capo-rctx->oldcapo-1;
415     }
416   }
417   /* skip the first and the last steps of each stride */
418   if (!s->userevolve && localstepnum==0) PetscFunctionReturn(0);
419 
420   /* checkpoint to memmory */
421   if (localstepnum==s->top) { /* overwrite the top checkpoint, this might happen when the time interval is split into several smaller ones, each corresponding to a call of TSSolve() */
422     ierr = StackTop(s,&e);CHKERRQ(ierr);
423     ierr = VecCopy(X,e->X);CHKERRQ(ierr);
424     ierr = TSGetStages(ts,&s->numY,&Y);CHKERRQ(ierr);
425     for (i=0;i<s->numY;i++) {
426       ierr = VecCopy(Y[i],e->Y[i]);CHKERRQ(ierr);
427     }
428     e->stepnum  = stepnum;
429     e->time     = time;
430     ierr        = TSGetPrevTime(ts,&timeprev);CHKERRQ(ierr);
431     e->timeprev = timeprev;
432   } else {
433     ierr = PetscCalloc1(1,&e);CHKERRQ(ierr);
434     ierr = VecDuplicate(X,&e->X);CHKERRQ(ierr);
435     ierr = VecCopy(X,e->X);CHKERRQ(ierr);
436     ierr = TSGetStages(ts,&s->numY,&Y);CHKERRQ(ierr);
437     ierr = VecDuplicateVecs(Y[0],s->numY,&e->Y);CHKERRQ(ierr);
438     for (i=0;i<s->numY;i++) {
439       ierr = VecCopy(Y[i],e->Y[i]);CHKERRQ(ierr);
440     }
441     e->stepnum  = stepnum;
442     e->time     = time;
443     /* for consistency */
444     if (stepnum == 0) {
445       e->timeprev = e->time - ts->time_step;
446     } else {
447       ierr = TSGetPrevTime(ts,&timeprev);CHKERRQ(ierr);
448       e->timeprev = timeprev;
449     }
450     ierr = StackPush(s,e);CHKERRQ(ierr);
451   }
452   PetscFunctionReturn(0);
453 }
454 
455 #undef __FUNCT__
456 #define __FUNCT__ "TSTrajectoryGet_Memory"
457 PetscErrorCode TSTrajectoryGet_Memory(TSTrajectory tj,TS ts,PetscInt stepnum,PetscReal *t)
458 {
459   Vec            *Y;
460   PetscInt       i;
461   StackElement   e;
462   Stack          *s = (Stack*)tj->data;
463   PetscReal      stepsize;
464   PetscInt       whattodo,rank,localstepnum,id,shift;
465   PetscErrorCode ierr;
466 
467   PetscFunctionBegin;
468   if (s->stride>1) { /* multilevel mode */
469     localstepnum = stepnum%s->stride;
470     if (localstepnum==0 && stepnum!=0 && stepnum!=s->total_steps) {
471       id = stepnum/s->stride;
472       ierr = StackLoadAll(ts,s,id);CHKERRQ(ierr);
473       //ierr = StackPop(s,&e);CHKERRQ(ierr); /* pop out stepnum 0 */
474       if (s->userevolve) {
475         s->top                  = s->max_cps-1;
476         s->rctx                 = s->rctx;
477         s->rctx->reverseonestep = PETSC_TRUE;
478         s->rctx->check          = s->top;
479       } else {
480         s->top = s->stride-2;
481       }
482     }
483   } else {
484     localstepnum = stepnum;
485   }
486   if (s->userevolve && s->rctx->reverseonestep) { /* ready for the reverse step */
487     ierr = TSGetTimeStep(ts,&stepsize);CHKERRQ(ierr);
488     ierr = TSSetTimeStep(ts,-stepsize);CHKERRQ(ierr);
489     s->rctx->reverseonestep = PETSC_FALSE;
490     PetscFunctionReturn(0);
491   }
492   if (stepnum==s->total_steps || localstepnum==0) {
493     ierr = TSGetTimeStep(ts,&stepsize);CHKERRQ(ierr);
494     ierr = TSSetTimeStep(ts,-stepsize);CHKERRQ(ierr);
495     PetscFunctionReturn(0);
496   }
497   /* restore a checkpoint */
498   ierr = StackTop(s,&e);CHKERRQ(ierr);
499   ierr = VecCopy(e->X,ts->vec_sol);CHKERRQ(ierr);
500   ierr = TSGetStages(ts,&s->numY,&Y);CHKERRQ(ierr);
501   for (i=0;i<s->numY;i++) {
502     ierr = VecCopy(e->Y[i],Y[i]);CHKERRQ(ierr);
503   }
504   *t = e->time;
505 
506   if (e->stepnum < stepnum) { /* need recomputation */
507     shift = stepnum-localstepnum;
508     s->rctx->capo = localstepnum;
509     whattodo = wrap_revolve(&s->rctx->check,&s->rctx->capo,&s->rctx->fine,&s->rctx->snaps_in,&s->rctx->info,&rank);
510 #ifdef TJ_VERBOSE
511     printwhattodo(whattodo,s->rctx,shift);
512 #endif
513     ierr = TSSetTimeStep(ts,(*t)-e->timeprev);CHKERRQ(ierr);
514     /* reset ts context */
515     PetscInt steps = ts->steps;
516     ts->steps      = e->stepnum;
517     ts->ptime      = e->time;
518     ts->ptime_prev = e->timeprev;
519     for (i=e->stepnum;i<stepnum;i++) { /* assume fixed step size */
520       ierr = TSTrajectorySet(ts->trajectory,ts,ts->steps,ts->ptime,ts->vec_sol);CHKERRQ(ierr);
521       ierr = TSMonitor(ts,ts->steps,ts->ptime,ts->vec_sol);CHKERRQ(ierr);
522       ierr = TSStep(ts);CHKERRQ(ierr);
523       if (ts->event) {
524         ierr = TSEventMonitor(ts);CHKERRQ(ierr);
525       }
526       if (!ts->steprollback) {
527         ierr = TSPostStep(ts);CHKERRQ(ierr);
528       }
529     }
530     /* reverseonestep must be true after the for loop */
531     ts->steps = steps;
532     ts->total_steps = stepnum;
533     ierr = TSGetTimeStep(ts,&stepsize);CHKERRQ(ierr);
534     ierr = TSSetTimeStep(ts,-stepsize);CHKERRQ(ierr);
535     if (stepnum-e->stepnum==1) { /* the restored checkpoint can be deleted now */
536       ierr = StackPop(s,&e);CHKERRQ(ierr);
537       ierr = VecDestroy(&e->X);CHKERRQ(ierr);
538       ierr = VecDestroyVecs(s->numY,&e->Y);CHKERRQ(ierr);
539       ierr = PetscFree(e);CHKERRQ(ierr);
540     }
541     s->rctx->reverseonestep = PETSC_FALSE;
542   } else if (e->stepnum == stepnum) { /* restore the data directly from checkpoints */
543     ierr = TSSetTimeStep(ts,-(*t)+e->timeprev);CHKERRQ(ierr);
544     ierr = StackPop(s,&e);CHKERRQ(ierr);
545     ierr = VecDestroy(&e->X);CHKERRQ(ierr);
546     ierr = VecDestroyVecs(s->numY,&e->Y);CHKERRQ(ierr);
547     ierr = PetscFree(e);CHKERRQ(ierr);
548   } else {
549     SETERRQ2(s->comm,PETSC_ERR_ARG_OUTOFRANGE,"The current step no. is %D, but the step number at top of the stack is %D",stepnum,e->stepnum);
550   }
551 
552   PetscFunctionReturn(0);
553 }
554 
555 #undef __FUNCT__
556 #define __FUNCT__ "TSTrajectoryDestroy_Memory"
557 PETSC_EXTERN PetscErrorCode TSTrajectoryDestroy_Memory(TSTrajectory tj)
558 {
559   Stack          *s = (Stack*)tj->data;
560   PetscErrorCode ierr;
561 
562   PetscFunctionBegin;
563   ierr = StackDestroy(s);CHKERRQ(ierr);
564   PetscFunctionReturn(0);
565 }
566 
567 /*MC
568       TSTRAJECTORYMEMORY - Stores each solution of the ODE/ADE in memory
569 
570   Level: intermediate
571 
572 .seealso:  TSTrajectoryCreate(), TS, TSTrajectorySetType()
573 
574 M*/
575 #undef __FUNCT__
576 #define __FUNCT__ "TSTrajectoryCreate_Memory"
577 PETSC_EXTERN PetscErrorCode TSTrajectoryCreate_Memory(TSTrajectory tj,TS ts)
578 {
579   Stack          *s;
580   PetscErrorCode ierr;
581 
582   PetscFunctionBegin;
583   tj->ops->set            = TSTrajectorySet_Memory;
584   tj->ops->get            = TSTrajectoryGet_Memory;
585   tj->ops->setup          = TSTrajectorySetUp_Memory;
586   tj->ops->destroy        = TSTrajectoryDestroy_Memory;
587   tj->ops->setfromoptions = TSTrajectorySetFromOptions_Memory;
588 
589   ierr = PetscCalloc1(1,&s);CHKERRQ(ierr);
590   s->max_cps = -1; /* -1 indicates that it is not set */
591   s->stride  = 0; /* if not zero, two-level checkpointing will be used */
592   tj->data   = s;
593   PetscFunctionReturn(0);
594 }
595