xref: /petsc/src/ts/impls/mimex/mimex.c (revision b02fdbf792e382cbf33467c854f6ffe5a15388ca)
1 /*
2        Code for Timestepping with my makeshift IMEX.
3 */
4 #include <petsc-private/tsimpl.h>                /*I   "petscts.h"   I*/
5 #include <petscds.h>
6 #include <petscdmplex.h>
7 
8 typedef struct {
9   Vec       Xdot, update;
10   PetscReal stage_time;
11 } TS_Mimex;
12 
13 #undef __FUNCT__
14 #define __FUNCT__ "TSMimexGetX0AndXdot"
15 static PetscErrorCode TSMimexGetX0AndXdot(TS ts, DM dm, Vec *X0, Vec *Xdot)
16 {
17   TS_Mimex      *mimex = (TS_Mimex *) ts->data;
18   PetscErrorCode ierr;
19 
20   PetscFunctionBegin;
21   if (X0) {
22     if (dm && dm != ts->dm) {ierr = DMGetNamedGlobalVector(dm, "TSMimex_X0", X0);CHKERRQ(ierr);}
23     else                    {*X0  = ts->vec_sol;}
24   }
25   if (Xdot) {
26     if (dm && dm != ts->dm) {ierr  = DMGetNamedGlobalVector(dm, "TSMimex_Xdot", Xdot);CHKERRQ(ierr);}
27     else                    {*Xdot = mimex->Xdot;}
28   }
29   PetscFunctionReturn(0);
30 }
31 
32 
33 #undef __FUNCT__
34 #define __FUNCT__ "TSMimexRestoreX0AndXdot"
35 static PetscErrorCode TSMimexRestoreX0AndXdot(TS ts, DM dm, Vec *X0, Vec *Xdot)
36 {
37   PetscErrorCode ierr;
38 
39   PetscFunctionBegin;
40   if (X0)   if (dm && dm != ts->dm) {ierr = DMRestoreNamedGlobalVector(dm, "TSMimex_X0", X0);CHKERRQ(ierr);}
41   if (Xdot) if (dm && dm != ts->dm) {ierr = DMRestoreNamedGlobalVector(dm, "TSMimex_Xdot", Xdot);CHKERRQ(ierr);}
42   PetscFunctionReturn(0);
43 }
44 
45 /*
46   This defines the nonlinear equation that is to be solved with SNES
47   G(U) = F[t0+dt, U, (U-U0)*shift] = 0
48 */
49 #undef __FUNCT__
50 #define __FUNCT__ "SNESTSFormFunction_Mimex"
51 static PetscErrorCode SNESTSFormFunction_Mimex(SNES snes, Vec x, Vec y, TS ts)
52 {
53   TS_Mimex      *mimex = (TS_Mimex *) ts->data;
54   DM             dm, dmsave;
55   Vec            X0, Xdot;
56   PetscReal      shift = 1./ts->time_step;
57   PetscErrorCode ierr;
58 
59   PetscFunctionBegin;
60   ierr = SNESGetDM(snes, &dm);CHKERRQ(ierr);
61   ierr = TSMimexGetX0AndXdot(ts, dm, &X0, &Xdot);CHKERRQ(ierr);
62   ierr = VecAXPBYPCZ(Xdot, -shift, shift, 0, X0, x);CHKERRQ(ierr);
63 
64   /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */
65   dmsave = ts->dm;
66   ts->dm = dm;
67   ierr   = TSComputeIFunction(ts, mimex->stage_time, x, Xdot, y, PETSC_TRUE);CHKERRQ(ierr);
68   ts->dm = dmsave;
69   ierr   = TSMimexRestoreX0AndXdot(ts, dm, &X0, &Xdot);CHKERRQ(ierr);
70   PetscFunctionReturn(0);
71 }
72 
73 #undef __FUNCT__
74 #define __FUNCT__ "SNESTSFormJacobian_Mimex"
75 static PetscErrorCode SNESTSFormJacobian_Mimex(SNES snes, Vec x, Mat A, Mat B, TS ts)
76 {
77   TS_Mimex      *mimex = (TS_Mimex *) ts->data;
78   DM             dm, dmsave;
79   Vec            Xdot;
80   PetscReal      shift = 1./ts->time_step;
81   PetscErrorCode ierr;
82 
83   PetscFunctionBegin;
84   /* th->Xdot has already been computed in SNESTSFormFunction_Mimex (SNES guarantees this) */
85   ierr = SNESGetDM(snes, &dm);CHKERRQ(ierr);
86   ierr = TSMimexGetX0AndXdot(ts, dm, NULL, &Xdot);CHKERRQ(ierr);
87 
88   /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */
89   dmsave = ts->dm;
90   ts->dm = dm;
91   ierr   = TSComputeIJacobian(ts, mimex->stage_time, x, Xdot, shift, A, B, PETSC_TRUE);CHKERRQ(ierr);
92   ts->dm = dmsave;
93   ierr   = TSMimexRestoreX0AndXdot(ts, dm, NULL, &Xdot);CHKERRQ(ierr);
94   PetscFunctionReturn(0);
95 }
96 
97 #undef __FUNCT__
98 #define __FUNCT__ "TSStep_Mimex"
99 static PetscErrorCode TSStep_Mimex(TS ts)
100 {
101   TS_Mimex          *mimex = (TS_Mimex *) ts->data;
102   DM                 dm;
103   PetscDS            prob;
104   PetscSection       s;
105   Vec                sol = ts->vec_sol, update = mimex->update;
106   const PetscScalar *aupdate;
107   PetscScalar       *asol, dt = ts->time_step;
108   PetscInt           Nf, f, pStart, pEnd, p;
109   PetscErrorCode     ierr;
110 
111   PetscFunctionBegin;
112   ierr = TSGetDM(ts, &dm);CHKERRQ(ierr);
113   ierr = DMGetDS(dm, &prob);CHKERRQ(ierr);
114   ierr = DMGetDefaultSection(dm, &s);CHKERRQ(ierr);
115   ierr = PetscDSGetNumFields(prob, &Nf);CHKERRQ(ierr);
116   ierr = PetscSectionGetChart(s, &pStart, &pEnd);CHKERRQ(ierr);
117   ierr = TSPreStep(ts);CHKERRQ(ierr);
118   ierr = TSPreStage(ts, ts->ptime);CHKERRQ(ierr);
119   /* Compute implicit update */
120   mimex->stage_time = ts->ptime + ts->time_step;
121   ierr = VecCopy(sol, update);CHKERRQ(ierr);
122   ierr = SNESSolve(ts->snes, NULL, update);CHKERRQ(ierr);
123   ierr = VecGetArrayRead(update, &aupdate);CHKERRQ(ierr);
124   ierr = VecGetArray(sol, &asol);CHKERRQ(ierr);
125   for (f = 0; f < Nf; ++f) {
126     PetscObject  obj;
127     PetscClassId id;
128     PetscInt     pdim, d;
129 
130     ierr = PetscDSGetDiscretization(prob, f, &obj);CHKERRQ(ierr);
131     ierr = PetscObjectGetClassId(obj, &id);CHKERRQ(ierr);
132     if (id == PETSCFV_CLASSID) continue;
133     for (p = pStart; p < pEnd; ++p) {
134       PetscScalar *au, *as;
135       PetscInt     fdof, fcdof, d;
136 
137       ierr = PetscSectionGetFieldDof(s, p, f, &fdof);CHKERRQ(ierr);
138       ierr = PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);CHKERRQ(ierr);
139       ierr = DMPlexPointGlobalFieldRead(dm, p, f, aupdate, &au);CHKERRQ(ierr);
140       ierr = DMPlexPointGlobalFieldRef(dm, p, f, asol, &as);CHKERRQ(ierr);
141       for (d = 0; d < fdof-fcdof; ++d) as[d] = au[d];
142     }
143   }
144   ierr = VecRestoreArrayRead(update, &aupdate);CHKERRQ(ierr);
145   ierr = VecRestoreArray(sol, &asol);CHKERRQ(ierr);
146   /* Compute explicit update */
147   ierr = TSComputeRHSFunction(ts, ts->ptime, sol, update);CHKERRQ(ierr);
148   ierr = VecGetArrayRead(update, &aupdate);CHKERRQ(ierr);
149   ierr = VecGetArray(sol, &asol);CHKERRQ(ierr);
150   for (f = 0; f < Nf; ++f) {
151     PetscObject  obj;
152     PetscClassId id;
153     PetscInt     pdim, d;
154 
155     ierr = PetscDSGetDiscretization(prob, f, &obj);CHKERRQ(ierr);
156     ierr = PetscObjectGetClassId(obj, &id);CHKERRQ(ierr);
157     if (id != PETSCFV_CLASSID) continue;
158     for (p = pStart; p < pEnd; ++p) {
159       PetscScalar *au, *as;
160       PetscInt     fdof, fcdof, d;
161 
162       ierr = PetscSectionGetFieldDof(s, p, f, &fdof);CHKERRQ(ierr);
163       ierr = PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);CHKERRQ(ierr);
164       ierr = DMPlexPointGlobalFieldRead(dm, p, f, aupdate, &au);CHKERRQ(ierr);
165       ierr = DMPlexPointGlobalFieldRef(dm, p, f, asol, &as);CHKERRQ(ierr);
166       for (d = 0; d < fdof-fcdof; ++d) as[d] += dt*au[d];
167     }
168   }
169   ierr = VecRestoreArrayRead(update, &aupdate);CHKERRQ(ierr);
170   ierr = VecRestoreArray(sol, &asol);CHKERRQ(ierr);
171   ierr = TSPostStage(ts, ts->ptime, 0, &sol);CHKERRQ(ierr);
172   ts->ptime += ts->time_step;
173   ts->steps++;
174   PetscFunctionReturn(0);
175 }
176 
177 /*------------------------------------------------------------*/
178 
179 #undef __FUNCT__
180 #define __FUNCT__ "TSSetUp_Mimex"
181 static PetscErrorCode TSSetUp_Mimex(TS ts)
182 {
183   TS_Mimex       *mimex = (TS_Mimex*)ts->data;
184   PetscErrorCode ierr;
185 
186   PetscFunctionBegin;
187   ierr = VecDuplicate(ts->vec_sol, &mimex->update);CHKERRQ(ierr);
188   ierr = VecDuplicate(ts->vec_sol, &mimex->Xdot);CHKERRQ(ierr);
189   PetscFunctionReturn(0);
190 }
191 
192 #undef __FUNCT__
193 #define __FUNCT__ "TSReset_Mimex"
194 static PetscErrorCode TSReset_Mimex(TS ts)
195 {
196   TS_Mimex       *mimex = (TS_Mimex*)ts->data;
197   PetscErrorCode ierr;
198 
199   PetscFunctionBegin;
200   ierr = VecDestroy(&mimex->update);CHKERRQ(ierr);
201   ierr = VecDestroy(&mimex->Xdot);CHKERRQ(ierr);
202   PetscFunctionReturn(0);
203 }
204 
205 #undef __FUNCT__
206 #define __FUNCT__ "TSDestroy_Mimex"
207 static PetscErrorCode TSDestroy_Mimex(TS ts)
208 {
209   PetscErrorCode ierr;
210 
211   PetscFunctionBegin;
212   ierr = TSReset_Mimex(ts);CHKERRQ(ierr);
213   ierr = PetscFree(ts->data);CHKERRQ(ierr);
214   PetscFunctionReturn(0);
215 }
216 /*------------------------------------------------------------*/
217 
218 #undef __FUNCT__
219 #define __FUNCT__ "TSSetFromOptions_Mimex"
220 static PetscErrorCode TSSetFromOptions_Mimex(PetscOptions *PetscOptionsObject,TS ts)
221 {
222   PetscFunctionBegin;
223   PetscFunctionReturn(0);
224 }
225 
226 #undef __FUNCT__
227 #define __FUNCT__ "TSView_Mimex"
228 static PetscErrorCode TSView_Mimex(TS ts,PetscViewer viewer)
229 {
230   PetscFunctionBegin;
231   PetscFunctionReturn(0);
232 }
233 
234 #undef __FUNCT__
235 #define __FUNCT__ "TSInterpolate_Mimex"
236 static PetscErrorCode TSInterpolate_Mimex(TS ts,PetscReal t,Vec X)
237 {
238   PetscReal      alpha = (ts->ptime - t)/ts->time_step;
239   PetscErrorCode ierr;
240 
241   PetscFunctionBegin;
242   ierr = VecAXPBY(ts->vec_sol,1.0-alpha,alpha,X);CHKERRQ(ierr);
243   PetscFunctionReturn(0);
244 }
245 
246 #undef __FUNCT__
247 #define __FUNCT__ "TSComputeLinearStability_Mimex"
248 PetscErrorCode TSComputeLinearStability_Mimex(TS ts,PetscReal xr,PetscReal xi,PetscReal *yr,PetscReal *yi)
249 {
250   PetscFunctionBegin;
251   *yr = 1.0 + xr;
252   *yi = xi;
253   PetscFunctionReturn(0);
254 }
255 /* ------------------------------------------------------------ */
256 
257 /*MC
258       TSMIMEX - ODE solver using the explicit forward Mimex method
259 
260   Level: beginner
261 
262 .seealso:  TSCreate(), TS, TSSetType(), TSBEULER
263 
264 M*/
265 #undef __FUNCT__
266 #define __FUNCT__ "TSCreate_Mimex"
267 PETSC_EXTERN PetscErrorCode TSCreate_Mimex(TS ts)
268 {
269   TS_Mimex       *mimex;
270   PetscErrorCode ierr;
271 
272   PetscFunctionBegin;
273   ts->ops->setup           = TSSetUp_Mimex;
274   ts->ops->step            = TSStep_Mimex;
275   ts->ops->reset           = TSReset_Mimex;
276   ts->ops->destroy         = TSDestroy_Mimex;
277   ts->ops->setfromoptions  = TSSetFromOptions_Mimex;
278   ts->ops->view            = TSView_Mimex;
279   ts->ops->interpolate     = TSInterpolate_Mimex;
280   ts->ops->linearstability = TSComputeLinearStability_Mimex;
281   ts->ops->snesfunction    = SNESTSFormFunction_Mimex;
282   ts->ops->snesjacobian    = SNESTSFormJacobian_Mimex;
283 
284   ierr = PetscNewLog(ts,&mimex);CHKERRQ(ierr);
285   ts->data = (void*)mimex;
286   PetscFunctionReturn(0);
287 }
288