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