Actual source code: mimex.c

petsc-3.6.1 2015-08-06
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:   TSPreStep(ts);
182:   TSPreStage(ts, ts->ptime);
183:   /* Compute implicit update */
184:   mimex->stage_time = ts->ptime + ts->time_step;
185:   VecCopy(sol, update);
186:   SNESSolve(ts->snes, NULL, update);
187:   VecGetArrayRead(update, &aupdate);
188:   VecGetArray(sol, &asol);
189:   for (f = 0; f < Nf; ++f) {
190:     PetscBool implicit;

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

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

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

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


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

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

262: static PetscErrorCode TSStep_Mimex(TS ts)
263: {
264:   TS_Mimex       *mimex = (TS_Mimex*)ts->data;
265:   PetscErrorCode  ierr;

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

279: /*------------------------------------------------------------*/

283: static PetscErrorCode TSSetUp_Mimex(TS ts)
284: {
285:   TS_Mimex       *mimex = (TS_Mimex*)ts->data;

289:   VecDuplicate(ts->vec_sol, &mimex->update);
290:   VecDuplicate(ts->vec_sol, &mimex->Xdot);
291:   return(0);
292: }

296: static PetscErrorCode TSReset_Mimex(TS ts)
297: {
298:   TS_Mimex       *mimex = (TS_Mimex*)ts->data;

302:   VecDestroy(&mimex->update);
303:   VecDestroy(&mimex->Xdot);
304:   return(0);
305: }

309: static PetscErrorCode TSDestroy_Mimex(TS ts)
310: {

314:   TSReset_Mimex(ts);
315:   PetscFree(ts->data);
316:   return(0);
317: }
318: /*------------------------------------------------------------*/

322: static PetscErrorCode TSSetFromOptions_Mimex(PetscOptions *PetscOptionsObject, TS ts)
323: {
324:   TS_Mimex      *mimex = (TS_Mimex *) ts->data;

328:   PetscOptionsHead(PetscOptionsObject, "MIMEX ODE solver options");
329:   {
330:     PetscOptionsInt("-ts_mimex_version", "Algorithm version", "TSMimexSetVersion", mimex->version, &mimex->version, NULL);
331:   }
332:   PetscOptionsTail();
333:   return(0);
334: }

338: static PetscErrorCode TSView_Mimex(TS ts,PetscViewer viewer)
339: {
340:   TS_Mimex      *mimex = (TS_Mimex *) ts->data;
341:   PetscBool      iascii;

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

355: static PetscErrorCode TSInterpolate_Mimex(TS ts,PetscReal t,Vec X)
356: {
357:   PetscReal      alpha = (ts->ptime - t)/ts->time_step;

361:   VecAXPBY(ts->vec_sol,1.0-alpha,alpha,X);
362:   return(0);
363: }

367: PetscErrorCode TSComputeLinearStability_Mimex(TS ts,PetscReal xr,PetscReal xi,PetscReal *yr,PetscReal *yi)
368: {
370:   *yr = 1.0 + xr;
371:   *yi = xi;
372:   return(0);
373: }
374: /* ------------------------------------------------------------ */

376: /*MC
377:       TSMIMEX - ODE solver using the explicit forward Mimex method

379:   Level: beginner

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

383: M*/
386: PETSC_EXTERN PetscErrorCode TSCreate_Mimex(TS ts)
387: {
388:   TS_Mimex       *mimex;

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

403:   PetscNewLog(ts,&mimex);
404:   ts->data = (void*)mimex;

406:   mimex->version = 1;
407:   return(0);
408: }