xref: /petsc/src/ts/trajectory/impls/memory/trajmemory.c (revision 20a6ff2bf3e68ed71e4fbaba49bff68b457f8022)
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   PetscReal time;
9   Vec       *Y;
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==s->top) { /* overwrite the top checkpoint */
108     ierr = StackTop(s,&e);
109     ierr = VecCopy(X,e->X);CHKERRQ(ierr);
110     ierr = TSGetStages(ts,&ns,&Y);CHKERRQ(ierr);
111     for (i=0;i<ns;i++) {
112       ierr = VecCopy(Y[i],e->Y[i]);CHKERRQ(ierr);
113     }
114     e->stepnum  = stepnum;
115     e->time     = time;
116     ierr        = TSGetPrevTime(ts,&timeprev);CHKERRQ(ierr);
117     e->timeprev = timeprev;
118   } else {
119     ierr = PetscCalloc1(1,&e);
120     ierr = VecDuplicate(X,&e->X);CHKERRQ(ierr);
121     ierr = VecCopy(X,e->X);CHKERRQ(ierr);
122     ierr = TSGetStages(ts,&ns,&Y);CHKERRQ(ierr);
123     ierr = VecDuplicateVecs(Y[0],ns,&e->Y);CHKERRQ(ierr);
124     for (i=0;i<ns;i++) {
125       ierr = VecCopy(Y[i],e->Y[i]);CHKERRQ(ierr);
126     }
127     e->stepnum  = stepnum;
128     e->time     = time;
129     ierr        = TSGetPrevTime(ts,&timeprev);CHKERRQ(ierr);
130     e->timeprev = timeprev;
131     ierr        = StackPush(s,e);CHKERRQ(ierr);
132   }
133   PetscFunctionReturn(0);
134 }
135 
136 #undef __FUNCT__
137 #define __FUNCT__ "TSTrajectoryGet_Memory"
138 PetscErrorCode TSTrajectoryGet_Memory(TSTrajectory tj,TS ts,PetscInt step,PetscReal *t)
139 {
140   Vec            Sol,*Y;
141   PetscInt       nr,i;
142   StackElement   e;
143   Stack          *s = (Stack*)tj->data;
144   PetscErrorCode ierr;
145 
146   PetscFunctionBegin;
147   ierr = StackPop(s,&e);CHKERRQ(ierr);
148   ierr = TSGetSolution(ts,&Sol);CHKERRQ(ierr);
149   ierr = VecCopy(e->X,Sol);CHKERRQ(ierr);
150   ierr = TSGetStages(ts,&nr,&Y);CHKERRQ(ierr);
151   for (i=0;i<nr ;i++) {
152     ierr = VecCopy(e->Y[i],Y[i]);CHKERRQ(ierr);
153   }
154   *t   = e->time;
155   ierr = TSSetTimeStep(ts,-(*t)+e->timeprev);CHKERRQ(ierr);
156 
157   ierr = VecDestroy(&e->X);CHKERRQ(ierr);
158   ierr = VecDestroyVecs(s->numY,&e->Y);CHKERRQ(ierr);
159   ierr = PetscFree(e);CHKERRQ(ierr);
160   PetscFunctionReturn(0);
161 }
162 
163 #undef __FUNCT__
164 #define __FUNCT__ "TSTrajectoryDestroy_Memory"
165 PETSC_EXTERN PetscErrorCode TSTrajectoryDestroy_Memory(TSTrajectory tj)
166 {
167   Stack          *s = (Stack*)tj->data;
168   PetscErrorCode ierr;
169 
170   PetscFunctionBegin;
171   ierr = StackDestroy(s);CHKERRQ(ierr);
172   PetscFunctionReturn(0);
173 }
174 
175 /*MC
176       TSTRAJECTORYMEMORY - Stores each solution of the ODE/ADE in memory
177 
178   Level: intermediate
179 
180 .seealso:  TSTrajectoryCreate(), TS, TSTrajectorySetType()
181 
182 M*/
183 #undef __FUNCT__
184 #define __FUNCT__ "TSTrajectoryCreate_Memory"
185 PETSC_EXTERN PetscErrorCode TSTrajectoryCreate_Memory(TSTrajectory tj,TS ts)
186 {
187   PetscInt       nr;
188   Stack          *s;
189   PetscErrorCode ierr;
190 
191   PetscFunctionBegin;
192   tj->ops->set     = TSTrajectorySet_Memory;
193   tj->ops->get     = TSTrajectoryGet_Memory;
194   tj->ops->destroy = TSTrajectoryDestroy_Memory;
195 
196   ierr = PetscCalloc1(1,&s);
197   ierr = TSGetStages(ts,&nr,PETSC_IGNORE);CHKERRQ(ierr);
198   ierr = StackCreate(PetscObjectComm((PetscObject)ts),s,ts->max_steps+1,nr);CHKERRQ(ierr);
199 
200   tj->data = s;
201   PetscFunctionReturn(0);
202 }
203