Actual source code: mimex.c

petsc-3.7.3 2016-08-01
Report Typos and Errors
  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>

  8: typedef struct {
  9:   Vec       Xdot, update;
 10:   PetscReal stage_time;
 11:   PetscInt  version;
 12: } TS_Mimex;

 16: static PetscErrorCode TSMimexGetX0AndXdot(TS ts, DM dm, Vec *X0, Vec *Xdot)
 17: {
 18:   TS_Mimex      *mimex = (TS_Mimex *) ts->data;

 22:   if (X0) {
 23:     if (dm && dm != ts->dm) {DMGetNamedGlobalVector(dm, "TSMimex_X0", X0);}
 24:     else                    {*X0  = ts->vec_sol;}
 25:   }
 26:   if (Xdot) {
 27:     if (dm && dm != ts->dm) {DMGetNamedGlobalVector(dm, "TSMimex_Xdot", Xdot);}
 28:     else                    {*Xdot = mimex->Xdot;}
 29:   }
 30:   return(0);
 31: }

 35: static PetscErrorCode TSMimexRestoreX0AndXdot(TS ts, DM dm, Vec *X0, Vec *Xdot)
 36: {

 40:   if (X0)   if (dm && dm != ts->dm) {DMRestoreNamedGlobalVector(dm, "TSMimex_X0", X0);}
 41:   if (Xdot) if (dm && dm != ts->dm) {DMRestoreNamedGlobalVector(dm, "TSMimex_Xdot", Xdot);}
 42:   return(0);
 43: }

 47: static PetscErrorCode TSMimexGetXstarAndG(TS ts, DM dm, Vec *Xstar, Vec *G)
 48: {

 52:   DMGetNamedGlobalVector(dm, "TSMimex_Xstar", Xstar);
 53:   DMGetNamedGlobalVector(dm, "TSMimex_G", G);
 54:   return(0);
 55: }

 59: static PetscErrorCode TSMimexRestoreXstarAndG(TS ts, DM dm, Vec *Xstar, Vec *G)
 60: {

 64:   DMRestoreNamedGlobalVector(dm, "TSMimex_Xstar", Xstar);
 65:   DMRestoreNamedGlobalVector(dm, "TSMimex_G", G);
 66:   return(0);
 67: }

 69: /*
 70:   This defines the nonlinear equation that is to be solved with SNES
 71:   G(U) = F[t0+dt, U, (U-U0)*shift] = 0
 72: */
 75: static PetscErrorCode SNESTSFormFunction_Mimex(SNES snes, Vec x, Vec y, TS ts)
 76: {
 77:   TS_Mimex      *mimex = (TS_Mimex *) ts->data;
 78:   DM             dm, dmsave;
 79:   Vec            X0, Xdot;
 80:   PetscReal      shift = 1./ts->time_step;

 84:   SNESGetDM(snes, &dm);
 85:   TSMimexGetX0AndXdot(ts, dm, &X0, &Xdot);
 86:   VecAXPBYPCZ(Xdot, -shift, shift, 0, X0, x);

 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:   TSComputeIFunction(ts, mimex->stage_time, x, Xdot, y, PETSC_TRUE);
 92:   if (mimex->version == 1) {
 93:     DM                 dm;
 94:     PetscDS            prob;
 95:     PetscSection       s;
 96:     Vec                Xstar = NULL, G = NULL;
 97:     const PetscScalar *ax;
 98:     PetscScalar       *axstar;
 99:     PetscInt           Nf, f, pStart, pEnd, p;

101:     TSGetDM(ts, &dm);
102:     DMGetDS(dm, &prob);
103:     DMGetDefaultSection(dm, &s);
104:     PetscDSGetNumFields(prob, &Nf);
105:     PetscSectionGetChart(s, &pStart, &pEnd);
106:     TSMimexGetXstarAndG(ts, dm, &Xstar, &G);
107:     VecCopy(X0, Xstar);
108:     VecGetArrayRead(x, &ax);
109:     VecGetArray(Xstar, &axstar);
110:     for (f = 0; f < Nf; ++f) {
111:       PetscBool implicit;

113:       PetscDSGetImplicit(prob, f, &implicit);
114:       if (!implicit) continue;
115:       for (p = pStart; p < pEnd; ++p) {
116:         PetscScalar *a, *axs;
117:         PetscInt     fdof, fcdof, d;

119:         PetscSectionGetFieldDof(s, p, f, &fdof);
120:         PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);
121:         DMPlexPointGlobalFieldRead(dm, p, f, ax, &a);
122:         DMPlexPointGlobalFieldRef(dm, p, f, axstar, &axs);
123:         for (d = 0; d < fdof-fcdof; ++d) axs[d] = a[d];
124:       }
125:     }
126:     VecRestoreArrayRead(x, &ax);
127:     VecRestoreArray(Xstar, &axstar);
128:     TSComputeRHSFunction(ts, ts->ptime, Xstar, G);
129:     VecAXPY(y, -1.0, G);
130:     TSMimexRestoreXstarAndG(ts, dm, &Xstar, &G);
131:   }
132:   ts->dm = dmsave;
133:   TSMimexRestoreX0AndXdot(ts, dm, &X0, &Xdot);
134:   return(0);
135: }

139: static PetscErrorCode SNESTSFormJacobian_Mimex(SNES snes, Vec x, Mat A, Mat B, TS ts)
140: {
141:   TS_Mimex      *mimex = (TS_Mimex *) ts->data;
142:   DM             dm, dmsave;
143:   Vec            Xdot;
144:   PetscReal      shift = 1./ts->time_step;

148:   /* th->Xdot has already been computed in SNESTSFormFunction_Mimex (SNES guarantees this) */
149:   SNESGetDM(snes, &dm);
150:   TSMimexGetX0AndXdot(ts, dm, NULL, &Xdot);

152:   /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */
153:   dmsave = ts->dm;
154:   ts->dm = dm;
155:   TSComputeIJacobian(ts, mimex->stage_time, x, Xdot, shift, A, B, PETSC_TRUE);
156:   ts->dm = dmsave;
157:   TSMimexRestoreX0AndXdot(ts, dm, NULL, &Xdot);
158:   return(0);
159: }

163: static PetscErrorCode TSStep_Mimex_Split(TS ts)
164: {
165:   TS_Mimex          *mimex = (TS_Mimex *) ts->data;
166:   DM                 dm;
167:   PetscDS            prob;
168:   PetscSection       s;
169:   Vec                sol = ts->vec_sol, update = mimex->update;
170:   const PetscScalar *aupdate;
171:   PetscScalar       *asol, dt = ts->time_step;
172:   PetscInt           Nf, f, pStart, pEnd, p;
173:   PetscErrorCode     ierr;

176:   TSGetDM(ts, &dm);
177:   DMGetDS(dm, &prob);
178:   DMGetDefaultSection(dm, &s);
179:   PetscDSGetNumFields(prob, &Nf);
180:   PetscSectionGetChart(s, &pStart, &pEnd);
181:   TSPreStage(ts, ts->ptime);
182:   /* Compute implicit update */
183:   mimex->stage_time = ts->ptime + ts->time_step;
184:   VecCopy(sol, update);
185:   SNESSolve(ts->snes, NULL, update);
186:   VecGetArrayRead(update, &aupdate);
187:   VecGetArray(sol, &asol);
188:   for (f = 0; f < Nf; ++f) {
189:     PetscBool implicit;

191:     PetscDSGetImplicit(prob, f, &implicit);
192:     if (!implicit) continue;
193:     for (p = pStart; p < pEnd; ++p) {
194:       PetscScalar *au, *as;
195:       PetscInt     fdof, fcdof, d;

197:       PetscSectionGetFieldDof(s, p, f, &fdof);
198:       PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);
199:       DMPlexPointGlobalFieldRead(dm, p, f, aupdate, &au);
200:       DMPlexPointGlobalFieldRef(dm, p, f, asol, &as);
201:       for (d = 0; d < fdof-fcdof; ++d) as[d] = au[d];
202:     }
203:   }
204:   VecRestoreArrayRead(update, &aupdate);
205:   VecRestoreArray(sol, &asol);
206:   /* Compute explicit update */
207:   TSComputeRHSFunction(ts, ts->ptime, sol, update);
208:   VecGetArrayRead(update, &aupdate);
209:   VecGetArray(sol, &asol);
210:   for (f = 0; f < Nf; ++f) {
211:     PetscBool implicit;

213:     PetscDSGetImplicit(prob, f, &implicit);
214:     if (implicit) continue;
215:     for (p = pStart; p < pEnd; ++p) {
216:       PetscScalar *au, *as;
217:       PetscInt     fdof, fcdof, d;

219:       PetscSectionGetFieldDof(s, p, f, &fdof);
220:       PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);
221:       DMPlexPointGlobalFieldRead(dm, p, f, aupdate, &au);
222:       DMPlexPointGlobalFieldRef(dm, p, f, asol, &as);
223:       for (d = 0; d < fdof-fcdof; ++d) as[d] += dt*au[d];
224:     }
225:   }
226:   VecRestoreArrayRead(update, &aupdate);
227:   VecRestoreArray(sol, &asol);
228:   TSPostStage(ts, ts->ptime, 0, &sol);
229:   ts->ptime += ts->time_step;
230:   return(0);
231: }


236: /* Evalute F at U and G at U0 for explicit fields and U for implicit fields */
237: static PetscErrorCode TSStep_Mimex_Implicit(TS ts)
238: {
239:   TS_Mimex      *mimex  = (TS_Mimex *) ts->data;
240:   Vec            sol    = ts->vec_sol;
241:   Vec            update = mimex->update;

245:   TSPreStage(ts, ts->ptime);
246:   /* Compute implicit update */
247:   mimex->stage_time = ts->ptime + ts->time_step;
248:   ts->ptime += ts->time_step;
249:   VecCopy(sol, update);
250:   SNESSolve(ts->snes, NULL, update);
251:   VecCopy(update, sol);
252:   TSPostStage(ts, ts->ptime, 0, &sol);
253:   return(0);
254: }

258: static PetscErrorCode TSStep_Mimex(TS ts)
259: {
260:   TS_Mimex       *mimex = (TS_Mimex*)ts->data;
261:   PetscErrorCode  ierr;

264:   switch(mimex->version) {
265:   case 0:
266:     TSStep_Mimex_Split(ts); break;
267:   case 1:
268:     TSStep_Mimex_Implicit(ts); break;
269:   default:
270:     SETERRQ1(PetscObjectComm((PetscObject) ts), PETSC_ERR_ARG_OUTOFRANGE, "Unknown MIMEX version %d", mimex->version);
271:   }
272:   return(0);
273: }

275: /*------------------------------------------------------------*/

279: static PetscErrorCode TSSetUp_Mimex(TS ts)
280: {
281:   TS_Mimex       *mimex = (TS_Mimex*)ts->data;

285:   VecDuplicate(ts->vec_sol, &mimex->update);
286:   VecDuplicate(ts->vec_sol, &mimex->Xdot);
287:   return(0);
288: }

292: static PetscErrorCode TSReset_Mimex(TS ts)
293: {
294:   TS_Mimex       *mimex = (TS_Mimex*)ts->data;

298:   VecDestroy(&mimex->update);
299:   VecDestroy(&mimex->Xdot);
300:   return(0);
301: }

305: static PetscErrorCode TSDestroy_Mimex(TS ts)
306: {

310:   TSReset_Mimex(ts);
311:   PetscFree(ts->data);
312:   return(0);
313: }
314: /*------------------------------------------------------------*/

318: static PetscErrorCode TSSetFromOptions_Mimex(PetscOptionItems *PetscOptionsObject, TS ts)
319: {
320:   TS_Mimex      *mimex = (TS_Mimex *) ts->data;

324:   PetscOptionsHead(PetscOptionsObject, "MIMEX ODE solver options");
325:   {
326:     PetscOptionsInt("-ts_mimex_version", "Algorithm version", "TSMimexSetVersion", mimex->version, &mimex->version, NULL);
327:   }
328:   PetscOptionsTail();
329:   return(0);
330: }

334: static PetscErrorCode TSView_Mimex(TS ts,PetscViewer viewer)
335: {
336:   TS_Mimex      *mimex = (TS_Mimex *) ts->data;
337:   PetscBool      iascii;

341:   PetscObjectTypeCompare((PetscObject) viewer, PETSCVIEWERASCII, &iascii);
342:   if (iascii) {
343:     PetscViewerASCIIPrintf(viewer, "  Version = %D\n", mimex->version);
344:   }
345:   if (ts->snes) {SNESView(ts->snes, viewer);}
346:   return(0);
347: }

351: static PetscErrorCode TSInterpolate_Mimex(TS ts,PetscReal t,Vec X)
352: {
353:   PetscReal      alpha = (ts->ptime - t)/ts->time_step;

357:   VecAXPBY(ts->vec_sol,1.0-alpha,alpha,X);
358:   return(0);
359: }

363: static PetscErrorCode TSComputeLinearStability_Mimex(TS ts,PetscReal xr,PetscReal xi,PetscReal *yr,PetscReal *yi)
364: {
366:   *yr = 1.0 + xr;
367:   *yi = xi;
368:   return(0);
369: }
370: /* ------------------------------------------------------------ */

372: /*MC
373:       TSMIMEX - ODE solver using the explicit forward Mimex method

375:   Level: beginner

377: .seealso:  TSCreate(), TS, TSSetType(), TSBEULER

379: M*/
382: PETSC_EXTERN PetscErrorCode TSCreate_Mimex(TS ts)
383: {
384:   TS_Mimex       *mimex;

388:   ts->ops->setup           = TSSetUp_Mimex;
389:   ts->ops->step            = TSStep_Mimex;
390:   ts->ops->reset           = TSReset_Mimex;
391:   ts->ops->destroy         = TSDestroy_Mimex;
392:   ts->ops->setfromoptions  = TSSetFromOptions_Mimex;
393:   ts->ops->view            = TSView_Mimex;
394:   ts->ops->interpolate     = TSInterpolate_Mimex;
395:   ts->ops->linearstability = TSComputeLinearStability_Mimex;
396:   ts->ops->snesfunction    = SNESTSFormFunction_Mimex;
397:   ts->ops->snesjacobian    = SNESTSFormJacobian_Mimex;

399:   PetscNewLog(ts,&mimex);
400:   ts->data = (void*)mimex;

402:   mimex->version = 1;
403:   return(0);
404: }