xref: /petsc/src/ts/trajectory/impls/memory/trajmemory.c (revision 74cbb51d9a6a64c8a11b14f35ec5d9dc9b556ec6)
1 #include <petsc/private/tsimpl.h>        /*I "petscts.h"  I*/
2 #include <petscsys.h>
3 
4 extern int wrap_revolve(int* check,int* capo,int* fine,int *snaps_in,int* info);
5 
6 typedef struct _StackElement {
7   PetscInt  stepnum;
8   Vec       X;
9   Vec       *Y;
10   PetscReal time;
11   PetscReal timeprev;
12 } *StackElement;
13 
14 typedef struct _Stack {
15    PetscBool    usecontroller;
16    PetscBool    reverseonestep;
17    PetscInt     stepsleft;
18    PetscInt     check;
19    PetscInt     capo;
20    PetscInt     fine;
21    PetscInt     info;
22    PetscInt     top;         /* The top of the stack */
23    PetscInt     maxelements; /* The maximum stack size */
24    PetscInt     numY;
25    MPI_Comm     comm;
26    StackElement *stack;      /* The storage */
27  } Stack;
28 
29 static PetscErrorCode StackCreate(MPI_Comm,Stack *,PetscInt,PetscInt);
30 static PetscErrorCode StackDestroy(Stack*);
31 static PetscErrorCode StackPush(Stack*,StackElement);
32 static PetscErrorCode StackPop(Stack*,StackElement*);
33 static PetscErrorCode StackTop(Stack*,StackElement*);
34 
35 #undef __FUNCT__
36 #define __FUNCT__ "StackCreate"
37 static PetscErrorCode StackCreate(MPI_Comm comm,Stack *s,PetscInt size,PetscInt ny)
38 {
39   PetscErrorCode ierr;
40 
41   PetscFunctionBegin;
42   s->top         = -1;
43   s->maxelements = size;
44   s->comm        = comm;
45   s->numY        = ny;
46 
47   ierr = PetscMalloc1(s->maxelements*sizeof(StackElement),&s->stack);CHKERRQ(ierr);
48   ierr = PetscMemzero(s->stack,s->maxelements*sizeof(StackElement));CHKERRQ(ierr);
49   PetscFunctionReturn(0);
50 }
51 
52 #undef __FUNCT__
53 #define __FUNCT__ "StackDestroy"
54 static PetscErrorCode StackDestroy(Stack *s)
55 {
56   PetscInt       i;
57   PetscErrorCode ierr;
58 
59   PetscFunctionBegin;
60   if (s->top>-1) {
61     for (i=0;i<=s->top;i++) {
62       ierr = VecDestroy(&s->stack[i]->X);CHKERRQ(ierr);
63       ierr = VecDestroyVecs(s->numY,&s->stack[i]->Y);CHKERRQ(ierr);
64       ierr = PetscFree(s->stack[i]);CHKERRQ(ierr);
65     }
66   }
67   ierr = PetscFree(s->stack);CHKERRQ(ierr);
68   ierr = PetscFree(s);CHKERRQ(ierr);
69   PetscFunctionReturn(0);
70 }
71 
72 #undef __FUNCT__
73 #define __FUNCT__ "StackPush"
74 static PetscErrorCode StackPush(Stack *s,StackElement e)
75 {
76   PetscFunctionBegin;
77   if (s->top+1 >= s->maxelements) SETERRQ1(s->comm,PETSC_ERR_MEMC,"Maximum stack size (%D) exceeded",s->maxelements);
78   s->stack[++s->top] = e;
79   PetscFunctionReturn(0);
80 }
81 
82 #undef __FUNCT__
83 #define __FUNCT__ "StackPop"
84 static PetscErrorCode StackPop(Stack *s,StackElement *e)
85 {
86   PetscFunctionBegin;
87   if (s->top == -1) SETERRQ(s->comm,PETSC_ERR_MEMC,"Emptry stack");
88   *e = s->stack[s->top--];
89   PetscFunctionReturn(0);
90 }
91 
92 #undef __FUNCT__
93 #define __FUNCT__ "StackTop"
94 static PetscErrorCode StackTop(Stack *s,StackElement *e)
95 {
96   PetscFunctionBegin;
97   *e = s->stack[s->top];
98   PetscFunctionReturn(0);
99 }
100 
101 #undef __FUNCT__
102 #define __FUNCT__ "TSTrajectorySet_Memory"
103 PetscErrorCode TSTrajectorySet_Memory(TSTrajectory tj,TS ts,PetscInt stepnum,PetscReal time,Vec X)
104 {
105   PetscInt       ns,i;
106   Vec            *Y;
107   PetscReal      timeprev;
108   StackElement   e;
109   Stack          *s = (Stack*)tj->data;
110   PetscInt       whattodo,oldcapo;
111   PetscErrorCode ierr;
112 
113   PetscFunctionBegin;
114   if (stepnum<s->top) SETERRQ(s->comm,PETSC_ERR_MEMC,"Illegal modification of a non-top stack element");
115 
116   if (s->usecontroller) {
117     printf("left %d\n",s->stepsleft);
118     if (s->reverseonestep) PetscFunctionReturn(0);
119     if (s->stepsleft==0) { /* let the controller determine what to do next */
120       s->capo = stepnum;
121       oldcapo = s->capo;
122       printf("check=%d,capo=%d,fine=%d,snaps_in=%d\t",s->check,s->capo,s->fine,s->maxelements);
123       whattodo = wrap_revolve(&s->check,&s->capo,&s->fine,&s->maxelements,&s->info);
124 switch(whattodo) {
125       case 1:
126         printf("Advance from %d to %d.\n", oldcapo, s->capo);
127         break;
128       case 2:
129         printf("Store in checkpoint number %d\n",s->check);
130         break;
131       case 3:
132         printf("First turn: Initialize adjoints and reverse first step.\n");
133         break;
134       case 4:
135         printf("Forward and reverse one step.\n");
136         break;
137       case 5:
138         printf("Restore in checkpoint number %d\n",s->check);
139         break;
140       case -1:
141         printf("Error!");
142         break;
143 }
144       if (whattodo==-1) SETERRQ(s->comm,PETSC_ERR_MEMC,"Error in the controller");
145       if (whattodo==1) {
146         s->stepsleft = s->capo-oldcapo-1;
147         PetscFunctionReturn(0); /* do not need to checkpoint */
148       }
149       if (whattodo==3 || whattodo==4) {
150         s->reverseonestep = PETSC_TRUE;
151         PetscFunctionReturn(0);
152       }
153       if (whattodo==5) {
154         printf("check=%d,capo=%d,fine=%d,snaps_in=%d\t",s->check,s->capo,s->fine,s->maxelements);
155         oldcapo = s->capo;
156         whattodo = wrap_revolve(&s->check,&s->capo,&s->fine,&s->maxelements,&s->info);
157         if (whattodo!=1) SETERRQ(s->comm,PETSC_ERR_MEMC,"Something weird happened");
158         s->stepsleft = s->capo-oldcapo;
159 switch(whattodo) {
160       case 1:
161         printf("5Advance from %d to %d.\n", oldcapo, s->capo);
162         break;
163       case 2:
164         printf("5Store in checkpoint number %d\n",s->check);
165         break;
166       case 3:
167         printf("5First turn: Initialize adjoints and reverse first step.\n");
168         break;
169       case 4:
170         printf("5Forward and reverse one step.\n");
171         break;
172       case 5:
173         printf("5Restore in checkpoint number %d\n",s->check);
174         break;
175       case -1:
176         printf("Error!");
177         break;
178 }
179         PetscFunctionReturn(0);
180       }
181       if (whattodo==2) {
182         printf("check=%d,capo=%d,fine=%d,snaps_in=%d\t",s->check,s->capo,s->fine,s->maxelements);
183         oldcapo = s->capo;
184         whattodo = wrap_revolve(&s->check,&s->capo,&s->fine,&s->maxelements,&s->info);
185         if (whattodo!=1) SETERRQ(s->comm,PETSC_ERR_MEMC,"Something weird happened");
186         if (stepnum==0) s->stepsleft = s->capo-oldcapo-1;
187         else s->stepsleft = s->capo-oldcapo-1;
188 switch(whattodo) {
189       case 1:
190         printf("2Advance from %d to %d.\n", oldcapo, s->capo);
191         break;
192       case 2:
193         printf("2Store in checkpoint number %d\n",s->check);
194         break;
195       case 3:
196         printf("2First turn: Initialize adjoints and reverse first step.\n");
197         break;
198       case 4:
199         printf("2Forward and reverse one step.\n");
200         break;
201       case 5:
202         printf("2Restore in checkpoint number %d\n",s->check);
203         break;
204       case -1:
205         printf("Error!");
206         break;
207 }
208       }
209     } else { /* advance s->stepsleft time steps without checkpointing */
210       s->stepsleft--;
211       PetscFunctionReturn(0);
212     }
213   }
214 
215   /* checkpoint to memmory */
216   if (stepnum==s->top) { /* overwrite the top checkpoint */
217     ierr = StackTop(s,&e);
218     ierr = VecCopy(X,e->X);CHKERRQ(ierr);
219     ierr = TSGetStages(ts,&ns,&Y);CHKERRQ(ierr);
220     for (i=0;i<ns;i++) {
221       ierr = VecCopy(Y[i],e->Y[i]);CHKERRQ(ierr);
222     }
223     e->stepnum  = stepnum;
224     e->time     = time;
225     ierr        = TSGetPrevTime(ts,&timeprev);CHKERRQ(ierr);
226     e->timeprev = timeprev;
227   } else {
228     ierr = PetscCalloc1(1,&e);
229     ierr = VecDuplicate(X,&e->X);CHKERRQ(ierr);
230     ierr = VecCopy(X,e->X);CHKERRQ(ierr);
231     ierr = TSGetStages(ts,&ns,&Y);CHKERRQ(ierr);
232     ierr = VecDuplicateVecs(Y[0],ns,&e->Y);CHKERRQ(ierr);
233     for (i=0;i<ns;i++) {
234       ierr = VecCopy(Y[i],e->Y[i]);CHKERRQ(ierr);
235     }
236     e->stepnum  = stepnum;
237     e->time     = time;
238     if (stepnum == 0) {
239       e->timeprev = e->time - ts->time_step; /* for consistency */
240     } else {
241       ierr        = TSGetPrevTime(ts,&timeprev);CHKERRQ(ierr);
242       e->timeprev = timeprev;
243     }
244     ierr        = StackPush(s,e);CHKERRQ(ierr);
245   }
246   PetscFunctionReturn(0);
247 }
248 
249 #undef __FUNCT__
250 #define __FUNCT__ "TSTrajectoryGet_Memory"
251 PetscErrorCode TSTrajectoryGet_Memory(TSTrajectory tj,TS ts,PetscInt stepnum,PetscReal *t)
252 {
253   Vec            *Y;
254   PetscInt       nr,i;
255   StackElement   e;
256   Stack          *s = (Stack*)tj->data;
257   PetscReal      stepsize;
258   PetscErrorCode ierr;
259 
260   PetscFunctionBegin;
261   if (s->usecontroller && s->reverseonestep) {
262     ierr = TSGetTimeStep(ts,&stepsize);CHKERRQ(ierr);
263     ierr = TSSetTimeStep(ts,-stepsize);CHKERRQ(ierr); /* go backward */
264     s->reverseonestep = PETSC_FALSE;
265     printf("reverse step %d\n",stepnum);
266     PetscFunctionReturn(0);
267   }
268 
269   /* restore a checkpoint */
270   ierr = StackTop(s,&e);CHKERRQ(ierr);
271   ierr = VecCopy(e->X,ts->vec_sol);CHKERRQ(ierr);
272   ierr = TSGetStages(ts,&nr,&Y);CHKERRQ(ierr);
273   for (i=0;i<nr ;i++) {
274     ierr = VecCopy(e->Y[i],Y[i]);CHKERRQ(ierr);
275   }
276   *t = e->time;
277 
278   printf("read stepnum=%d to get %d\n",e->stepnum,stepnum);
279   if (e->stepnum < stepnum) { /* need recomputation */
280     PetscInt whattodo;
281     s->capo = stepnum;
282     printf("check=%d,capo=%d,fine=%d,snaps_in=%d\t",s->check,s->capo,s->fine,s->maxelements);
283     whattodo = wrap_revolve(&s->check,&s->capo,&s->fine,&s->maxelements,&s->info);
284 switch(whattodo) {
285       case 5:
286         printf("R:Restore in checkpoint number %d\n",s->check);
287         break;
288       case -1:
289         printf("Error!");
290         break;
291 }
292     ierr = TSSetTimeStep(ts,(*t)-e->timeprev);CHKERRQ(ierr);
293     /* reset ts context */
294     PetscInt steps = ts->steps;
295     ts->steps      = e->stepnum;
296     ts->ptime      = e->time;
297     ts->ptime_prev = e->timeprev;
298     printf("need to recompute %d steps\n",stepnum-e->stepnum);
299     for (i=e->stepnum;i<stepnum;i++) { /* assume fixed step size */
300       ierr = TSTrajectorySet(ts->trajectory,ts,ts->steps,ts->ptime,ts->vec_sol);CHKERRQ(ierr);
301       ierr = TSMonitor(ts,ts->steps,ts->ptime,ts->vec_sol);CHKERRQ(ierr);
302       ierr = TSStep(ts);CHKERRQ(ierr);
303       if (ts->event) {
304         ierr = TSEventMonitor(ts);CHKERRQ(ierr);
305       }
306       if (!ts->steprollback) {
307         ierr = TSPostStep(ts);CHKERRQ(ierr);
308       }
309     }
310     if (!s->reverseonestep) printf("error\n");
311     ts->steps = steps;
312     ts->total_steps = stepnum;
313     ierr = TSGetTimeStep(ts,&stepsize);CHKERRQ(ierr);
314     ierr = TSSetTimeStep(ts,-stepsize);CHKERRQ(ierr); /* go backward */
315     if (stepnum-e->stepnum==1) {
316       ierr = StackPop(s,&e);CHKERRQ(ierr);
317       ierr = VecDestroy(&e->X);CHKERRQ(ierr);
318       ierr = VecDestroyVecs(s->numY,&e->Y);CHKERRQ(ierr);
319       ierr = PetscFree(e);CHKERRQ(ierr);
320     }
321     s->reverseonestep = PETSC_FALSE;
322     printf("reverse step %d\n",stepnum);
323   } else if (e->stepnum == stepnum) {
324     ierr = TSSetTimeStep(ts,-(*t)+e->timeprev);CHKERRQ(ierr); /* go backward */
325     ierr = StackPop(s,&e);CHKERRQ(ierr);
326     ierr = VecDestroy(&e->X);CHKERRQ(ierr);
327     ierr = VecDestroyVecs(s->numY,&e->Y);CHKERRQ(ierr);
328     ierr = PetscFree(e);CHKERRQ(ierr);
329   } else {
330     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);
331   }
332 
333   PetscFunctionReturn(0);
334 }
335 
336 #undef __FUNCT__
337 #define __FUNCT__ "TSTrajectoryDestroy_Memory"
338 PETSC_EXTERN PetscErrorCode TSTrajectoryDestroy_Memory(TSTrajectory tj)
339 {
340   Stack          *s = (Stack*)tj->data;
341   PetscErrorCode ierr;
342 
343   PetscFunctionBegin;
344   ierr = StackDestroy(s);CHKERRQ(ierr);
345   PetscFunctionReturn(0);
346 }
347 
348 /*MC
349       TSTRAJECTORYMEMORY - Stores each solution of the ODE/ADE in memory
350 
351   Level: intermediate
352 
353 .seealso:  TSTrajectoryCreate(), TS, TSTrajectorySetType()
354 
355 M*/
356 #undef __FUNCT__
357 #define __FUNCT__ "TSTrajectoryCreate_Memory"
358 PETSC_EXTERN PetscErrorCode TSTrajectoryCreate_Memory(TSTrajectory tj,TS ts)
359 {
360   PetscInt       nr,maxsteps;
361   Stack          *s;
362   PetscErrorCode ierr;
363 
364   PetscFunctionBegin;
365   tj->ops->set     = TSTrajectorySet_Memory;
366   tj->ops->get     = TSTrajectoryGet_Memory;
367   tj->ops->destroy = TSTrajectoryDestroy_Memory;
368 
369   ierr = PetscCalloc1(1,&s);
370   s->maxelements = 3; /* will be provided by users */
371   ierr = TSGetStages(ts,&nr,PETSC_IGNORE);CHKERRQ(ierr);
372 
373   maxsteps = PetscMin(ts->max_steps,(PetscInt)(ceil(ts->max_time/ts->time_step)));
374   if (s->maxelements-1<maxsteps) { /* Need to use a controller */
375     s->usecontroller  = PETSC_TRUE;
376     s->reverseonestep = PETSC_FALSE;
377     s->check          = -1;
378     s->capo           = 0;
379     s->fine           = maxsteps;
380     s->info           = 2;
381     ierr = StackCreate(PetscObjectComm((PetscObject)ts),s,s->maxelements,nr);CHKERRQ(ierr);
382   } else { /* Enough space for checkpointing all time steps */
383     s->usecontroller = PETSC_FALSE;
384     ierr = StackCreate(PetscObjectComm((PetscObject)ts),s,ts->max_steps+1,nr);CHKERRQ(ierr);
385   }
386   tj->data = s;
387   PetscFunctionReturn(0);
388 }
389