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