xref: /petsc/src/ts/trajectory/impls/memory/trajmemory.c (revision 93d756bd0e391448dbe9eb50dcdd4ba27619c510)
1 
2 #include <petsc/private/tsimpl.h>        /*I "petscts.h"  I*/
3 #include <petscsys.h>
4 
5 typedef struct _StackElement {
6   PetscInt  stepnum;
7   Vec       X;
8   Vec       *Y;
9   PetscReal time;
10   PetscReal timeprev;
11 } *StackElement;
12 
13 typedef struct _Stack {
14    PetscInt     top;         /* The top of the stack */
15    PetscInt     maxelements; /* The maximum stack size */
16    PetscInt     numY;
17    MPI_Comm     comm;
18    StackElement *stack;      /* The storage */
19  } Stack;
20 
21 static PetscErrorCode StackCreate(MPI_Comm,Stack *,PetscInt,PetscInt);
22 static PetscErrorCode StackDestroy(Stack*);
23 static PetscErrorCode StackPush(Stack*,StackElement);
24 static PetscErrorCode StackPop(Stack*,StackElement*);
25 static PetscErrorCode StackTop(Stack*,StackElement*);
26 
27 #undef __FUNCT__
28 #define __FUNCT__ "StackCreate"
29 static PetscErrorCode StackCreate(MPI_Comm comm,Stack *s,PetscInt size,PetscInt ny)
30 {
31   PetscErrorCode ierr;
32 
33   PetscFunctionBegin;
34   s->top         = -1;
35   s->maxelements = size;
36   s->comm        = comm;
37   s->numY        = ny;
38 
39   ierr = PetscMalloc1(s->maxelements*sizeof(StackElement),&s->stack);CHKERRQ(ierr);
40   ierr = PetscMemzero(s->stack,s->maxelements*sizeof(StackElement));CHKERRQ(ierr);
41   PetscFunctionReturn(0);
42 }
43 
44 #undef __FUNCT__
45 #define __FUNCT__ "StackDestroy"
46 static PetscErrorCode StackDestroy(Stack *s)
47 {
48   PetscInt       i;
49   PetscErrorCode ierr;
50 
51   PetscFunctionBegin;
52   if (s->top>-1) {
53     for (i=0;i<=s->top;i++) {
54       ierr = VecDestroy(&s->stack[i]->X);CHKERRQ(ierr);
55       ierr = VecDestroyVecs(s->numY,&s->stack[i]->Y);CHKERRQ(ierr);
56       ierr = PetscFree(s->stack[i]);CHKERRQ(ierr);
57     }
58   }
59   ierr = PetscFree(s->stack);CHKERRQ(ierr);
60   ierr = PetscFree(s);CHKERRQ(ierr);
61   PetscFunctionReturn(0);
62 }
63 
64 #undef __FUNCT__
65 #define __FUNCT__ "StackPush"
66 static PetscErrorCode StackPush(Stack *s,StackElement e)
67 {
68   PetscFunctionBegin;
69   if (s->top+1 >= s->maxelements) SETERRQ1(s->comm,PETSC_ERR_MEMC,"Maximum stack size (%D) exceeded",s->maxelements);
70   s->stack[++s->top] = e;
71   PetscFunctionReturn(0);
72 }
73 
74 #undef __FUNCT__
75 #define __FUNCT__ "StackPop"
76 static PetscErrorCode StackPop(Stack *s,StackElement *e)
77 {
78   PetscFunctionBegin;
79   if (s->top == -1) SETERRQ(s->comm,PETSC_ERR_MEMC,"Emptry stack");
80   *e = s->stack[s->top--];
81   PetscFunctionReturn(0);
82 }
83 
84 #undef __FUNCT__
85 #define __FUNCT__ "StackTop"
86 static PetscErrorCode StackTop(Stack *s,StackElement *e)
87 {
88   PetscFunctionBegin;
89   *e = s->stack[s->top];
90   PetscFunctionReturn(0);
91 }
92 
93 #undef __FUNCT__
94 #define __FUNCT__ "TSTrajectorySet_Memory"
95 PetscErrorCode TSTrajectorySet_Memory(TSTrajectory tj,TS ts,PetscInt stepnum,PetscReal time,Vec X)
96 {
97   PetscInt       ns,i;
98   Vec            *Y;
99   PetscReal      timeprev;
100   StackElement   e;
101   Stack          *s = (Stack*)tj->data;
102   PetscErrorCode ierr;
103 
104   PetscFunctionBegin;
105   if (stepnum<s->top) SETERRQ(s->comm,PETSC_ERR_MEMC,"Illegal modification of a non-top stack element");
106 
107   if (stepnum > (rand()%ts->max_steps)) PetscFunctionReturn(0);/* choose check points randomly */
108 
109   if (stepnum==s->top) { /* overwrite the top checkpoint */
110     ierr = StackTop(s,&e);
111     ierr = VecCopy(X,e->X);CHKERRQ(ierr);
112     ierr = TSGetStages(ts,&ns,&Y);CHKERRQ(ierr);
113     for (i=0;i<ns;i++) {
114       ierr = VecCopy(Y[i],e->Y[i]);CHKERRQ(ierr);
115     }
116     e->stepnum  = stepnum;
117     e->time     = time;
118     ierr        = TSGetPrevTime(ts,&timeprev);CHKERRQ(ierr);
119     e->timeprev = timeprev;
120   } else {
121     ierr = PetscCalloc1(1,&e);
122     ierr = VecDuplicate(X,&e->X);CHKERRQ(ierr);
123     ierr = VecCopy(X,e->X);CHKERRQ(ierr);
124     ierr = TSGetStages(ts,&ns,&Y);CHKERRQ(ierr);
125     ierr = VecDuplicateVecs(Y[0],ns,&e->Y);CHKERRQ(ierr);
126     for (i=0;i<ns;i++) {
127       ierr = VecCopy(Y[i],e->Y[i]);CHKERRQ(ierr);
128     }
129     e->stepnum  = stepnum;
130     e->time     = time;
131     if (stepnum == 0) {
132       e->timeprev = e->time - ts->time_step; /* for consistency */
133     } else {
134       ierr        = TSGetPrevTime(ts,&timeprev);CHKERRQ(ierr);
135       e->timeprev = timeprev;
136     }
137     ierr        = StackPush(s,e);CHKERRQ(ierr);
138   }
139   PetscFunctionReturn(0);
140 }
141 
142 #undef __FUNCT__
143 #define __FUNCT__ "TSTrajectoryGet_Memory"
144 PetscErrorCode TSTrajectoryGet_Memory(TSTrajectory tj,TS ts,PetscInt stepnum,PetscReal *t)
145 {
146   Vec            *Y;
147   PetscInt       nr,i;
148   StackElement   e;
149   Stack          *s = (Stack*)tj->data;
150   PetscErrorCode ierr;
151 
152   PetscFunctionBegin;
153 
154   ierr = StackTop(s,&e);CHKERRQ(ierr);
155   ierr = VecCopy(e->X,ts->vec_sol);CHKERRQ(ierr);
156   ierr = TSGetStages(ts,&nr,&Y);CHKERRQ(ierr);
157   for (i=0;i<nr ;i++) {
158     ierr = VecCopy(e->Y[i],Y[i]);CHKERRQ(ierr);
159   }
160   *t = e->time;
161 
162   if (e->stepnum < stepnum) { /* go forward */
163     ierr = TSSetTimeStep(ts,(*t)-e->timeprev);CHKERRQ(ierr);
164     /* reset ts context */
165     PetscInt steps = ts->steps;
166     ts->steps      = e->stepnum;
167     ts->ptime      = e->time;
168     ts->ptime_prev = e->timeprev;
169 
170     /* recompute to the specified step */
171     for (i=e->stepnum;i<stepnum;i++) { /* assume fixed step size */
172       ierr = TSMonitor(ts,ts->steps,ts->ptime,ts->vec_sol);CHKERRQ(ierr);
173       ierr = TSStep(ts);CHKERRQ(ierr);
174       if (ts->event) {
175         ierr = TSEventMonitor(ts);CHKERRQ(ierr);
176       }
177       if (!ts->steprollback) {
178         if (i+1<stepnum) { /* skip checkpointing at the last step */
179           ierr = TSTrajectorySet(ts->trajectory,ts,ts->steps,ts->ptime,ts->vec_sol);CHKERRQ(ierr);
180         }
181         ierr = TSPostStep(ts);CHKERRQ(ierr);
182       }
183     }
184     ts->steps = steps;
185     ts->total_steps = stepnum;
186     PetscReal stepsize;
187     ierr = TSGetTimeStep(ts,&stepsize);CHKERRQ(ierr);
188     ierr = TSSetTimeStep(ts,-stepsize);CHKERRQ(ierr); /* go backward */
189   } else if (e->stepnum == stepnum) {
190     ierr = TSSetTimeStep(ts,-(*t)+e->timeprev);CHKERRQ(ierr); /* go backward */
191     ierr = StackPop(s,&e);CHKERRQ(ierr);
192     ierr = VecDestroy(&e->X);CHKERRQ(ierr);
193     ierr = VecDestroyVecs(s->numY,&e->Y);CHKERRQ(ierr);
194     ierr = PetscFree(e);CHKERRQ(ierr);
195   } else {
196     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);
197   }
198 
199   PetscFunctionReturn(0);
200 }
201 
202 #undef __FUNCT__
203 #define __FUNCT__ "TSTrajectoryDestroy_Memory"
204 PETSC_EXTERN PetscErrorCode TSTrajectoryDestroy_Memory(TSTrajectory tj)
205 {
206   Stack          *s = (Stack*)tj->data;
207   PetscErrorCode ierr;
208 
209   PetscFunctionBegin;
210   ierr = StackDestroy(s);CHKERRQ(ierr);
211   PetscFunctionReturn(0);
212 }
213 
214 /*MC
215       TSTRAJECTORYMEMORY - Stores each solution of the ODE/ADE in memory
216 
217   Level: intermediate
218 
219 .seealso:  TSTrajectoryCreate(), TS, TSTrajectorySetType()
220 
221 M*/
222 #undef __FUNCT__
223 #define __FUNCT__ "TSTrajectoryCreate_Memory"
224 PETSC_EXTERN PetscErrorCode TSTrajectoryCreate_Memory(TSTrajectory tj,TS ts)
225 {
226   PetscInt       nr;
227   Stack          *s;
228   PetscErrorCode ierr;
229 
230   PetscFunctionBegin;
231   tj->ops->set     = TSTrajectorySet_Memory;
232   tj->ops->get     = TSTrajectoryGet_Memory;
233   tj->ops->destroy = TSTrajectoryDestroy_Memory;
234 
235   ierr = PetscCalloc1(1,&s);
236   ierr = TSGetStages(ts,&nr,PETSC_IGNORE);CHKERRQ(ierr);
237   ierr = StackCreate(PetscObjectComm((PetscObject)ts),s,ts->max_steps+1,nr);CHKERRQ(ierr);
238 
239   tj->data = s;
240   PetscFunctionReturn(0);
241 }
242