Actual source code: mimex.c

petsc-3.12.5 2020-03-29
Report Typos and Errors
  1: /*
  2:        Code for Timestepping with my makeshift IMEX.
  3: */
  4:  #include <petsc/private/tsimpl.h>
  5:  #include <petscds.h>
  6:  #include <petscsection.h>
  7:  #include <petscdmplex.h>

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

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

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

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

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

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

 47:   DMGetNamedGlobalVector(dm, "TSMimex_Xstar", Xstar);
 48:   DMGetNamedGlobalVector(dm, "TSMimex_G", G);
 49:   return(0);
 50: }

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

 57:   DMRestoreNamedGlobalVector(dm, "TSMimex_Xstar", Xstar);
 58:   DMRestoreNamedGlobalVector(dm, "TSMimex_G", G);
 59:   return(0);
 60: }

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

 75:   SNESGetDM(snes, &dm);
 76:   TSMimexGetX0AndXdot(ts, dm, &X0, &Xdot);
 77:   VecAXPBYPCZ(Xdot, -shift, shift, 0, X0, x);

 79:   /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */
 80:   dmsave = ts->dm;
 81:   ts->dm = dm;
 82:   TSComputeIFunction(ts, mimex->stage_time, x, Xdot, y, PETSC_TRUE);
 83:   if (mimex->version == 1) {
 84:     DM                 dm;
 85:     PetscDS            prob;
 86:     PetscSection       s;
 87:     Vec                Xstar = NULL, G = NULL;
 88:     const PetscScalar *ax;
 89:     PetscScalar       *axstar;
 90:     PetscInt           Nf, f, pStart, pEnd, p;

 92:     TSGetDM(ts, &dm);
 93:     DMGetDS(dm, &prob);
 94:     DMGetLocalSection(dm, &s);
 95:     PetscDSGetNumFields(prob, &Nf);
 96:     PetscSectionGetChart(s, &pStart, &pEnd);
 97:     TSMimexGetXstarAndG(ts, dm, &Xstar, &G);
 98:     VecCopy(X0, Xstar);
 99:     VecGetArrayRead(x, &ax);
100:     VecGetArray(Xstar, &axstar);
101:     for (f = 0; f < Nf; ++f) {
102:       PetscBool implicit;

104:       PetscDSGetImplicit(prob, f, &implicit);
105:       if (!implicit) continue;
106:       for (p = pStart; p < pEnd; ++p) {
107:         PetscScalar *a, *axs;
108:         PetscInt     fdof, fcdof, d;

110:         PetscSectionGetFieldDof(s, p, f, &fdof);
111:         PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);
112:         DMPlexPointGlobalFieldRead(dm, p, f, ax, &a);
113:         DMPlexPointGlobalFieldRef(dm, p, f, axstar, &axs);
114:         for (d = 0; d < fdof-fcdof; ++d) axs[d] = a[d];
115:       }
116:     }
117:     VecRestoreArrayRead(x, &ax);
118:     VecRestoreArray(Xstar, &axstar);
119:     TSComputeRHSFunction(ts, ts->ptime, Xstar, G);
120:     VecAXPY(y, -1.0, G);
121:     TSMimexRestoreXstarAndG(ts, dm, &Xstar, &G);
122:   }
123:   ts->dm = dmsave;
124:   TSMimexRestoreX0AndXdot(ts, dm, &X0, &Xdot);
125:   return(0);
126: }

128: static PetscErrorCode SNESTSFormJacobian_Mimex(SNES snes, Vec x, Mat A, Mat B, TS ts)
129: {
130:   TS_Mimex      *mimex = (TS_Mimex *) ts->data;
131:   DM             dm, dmsave;
132:   Vec            Xdot;
133:   PetscReal      shift = 1./ts->time_step;

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

141:   /* DM monkey-business allows user code to call TSGetDM() inside of functions evaluated on levels of FAS */
142:   dmsave = ts->dm;
143:   ts->dm = dm;
144:   TSComputeIJacobian(ts, mimex->stage_time, x, Xdot, shift, A, B, PETSC_TRUE);
145:   ts->dm = dmsave;
146:   TSMimexRestoreX0AndXdot(ts, dm, NULL, &Xdot);
147:   return(0);
148: }

150: static PetscErrorCode TSStep_Mimex_Split(TS ts)
151: {
152:   TS_Mimex          *mimex = (TS_Mimex *) ts->data;
153:   DM                 dm;
154:   PetscDS            prob;
155:   PetscSection       s;
156:   Vec                sol = ts->vec_sol, update = mimex->update;
157:   const PetscScalar *aupdate;
158:   PetscScalar       *asol, dt = ts->time_step;
159:   PetscInt           Nf, f, pStart, pEnd, p;
160:   PetscErrorCode     ierr;

163:   TSGetDM(ts, &dm);
164:   DMGetDS(dm, &prob);
165:   DMGetLocalSection(dm, &s);
166:   PetscDSGetNumFields(prob, &Nf);
167:   PetscSectionGetChart(s, &pStart, &pEnd);
168:   TSPreStage(ts, ts->ptime);
169:   /* Compute implicit update */
170:   mimex->stage_time = ts->ptime + ts->time_step;
171:   VecCopy(sol, update);
172:   SNESSolve(ts->snes, NULL, update);
173:   VecGetArrayRead(update, &aupdate);
174:   VecGetArray(sol, &asol);
175:   for (f = 0; f < Nf; ++f) {
176:     PetscBool implicit;

178:     PetscDSGetImplicit(prob, f, &implicit);
179:     if (!implicit) continue;
180:     for (p = pStart; p < pEnd; ++p) {
181:       PetscScalar *au, *as;
182:       PetscInt     fdof, fcdof, d;

184:       PetscSectionGetFieldDof(s, p, f, &fdof);
185:       PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);
186:       DMPlexPointGlobalFieldRead(dm, p, f, aupdate, &au);
187:       DMPlexPointGlobalFieldRef(dm, p, f, asol, &as);
188:       for (d = 0; d < fdof-fcdof; ++d) as[d] = au[d];
189:     }
190:   }
191:   VecRestoreArrayRead(update, &aupdate);
192:   VecRestoreArray(sol, &asol);
193:   /* Compute explicit update */
194:   TSComputeRHSFunction(ts, ts->ptime, sol, update);
195:   VecGetArrayRead(update, &aupdate);
196:   VecGetArray(sol, &asol);
197:   for (f = 0; f < Nf; ++f) {
198:     PetscBool implicit;

200:     PetscDSGetImplicit(prob, f, &implicit);
201:     if (implicit) continue;
202:     for (p = pStart; p < pEnd; ++p) {
203:       PetscScalar *au, *as;
204:       PetscInt     fdof, fcdof, d;

206:       PetscSectionGetFieldDof(s, p, f, &fdof);
207:       PetscSectionGetFieldConstraintDof(s, p, f, &fcdof);
208:       DMPlexPointGlobalFieldRead(dm, p, f, aupdate, &au);
209:       DMPlexPointGlobalFieldRef(dm, p, f, asol, &as);
210:       for (d = 0; d < fdof-fcdof; ++d) as[d] += dt*au[d];
211:     }
212:   }
213:   VecRestoreArrayRead(update, &aupdate);
214:   VecRestoreArray(sol, &asol);
215:   TSPostStage(ts, ts->ptime, 0, &sol);
216:   ts->ptime += ts->time_step;
217:   return(0);
218: }


221: /* Evalute F at U and G at U0 for explicit fields and U for implicit fields */
222: static PetscErrorCode TSStep_Mimex_Implicit(TS ts)
223: {
224:   TS_Mimex      *mimex  = (TS_Mimex *) ts->data;
225:   Vec            sol    = ts->vec_sol;
226:   Vec            update = mimex->update;

230:   TSPreStage(ts, ts->ptime);
231:   /* Compute implicit update */
232:   mimex->stage_time = ts->ptime + ts->time_step;
233:   ts->ptime += ts->time_step;
234:   VecCopy(sol, update);
235:   SNESSolve(ts->snes, NULL, update);
236:   VecCopy(update, sol);
237:   TSPostStage(ts, ts->ptime, 0, &sol);
238:   return(0);
239: }

241: static PetscErrorCode TSStep_Mimex(TS ts)
242: {
243:   TS_Mimex       *mimex = (TS_Mimex*)ts->data;
244:   PetscErrorCode  ierr;

247:   switch(mimex->version) {
248:   case 0:
249:     TSStep_Mimex_Split(ts); break;
250:   case 1:
251:     TSStep_Mimex_Implicit(ts); break;
252:   default:
253:     SETERRQ1(PetscObjectComm((PetscObject) ts), PETSC_ERR_ARG_OUTOFRANGE, "Unknown MIMEX version %d", mimex->version);
254:   }
255:   return(0);
256: }

258: /*------------------------------------------------------------*/

260: static PetscErrorCode TSSetUp_Mimex(TS ts)
261: {
262:   TS_Mimex       *mimex = (TS_Mimex*)ts->data;

266:   VecDuplicate(ts->vec_sol, &mimex->update);
267:   VecDuplicate(ts->vec_sol, &mimex->Xdot);
268:   return(0);
269: }

271: static PetscErrorCode TSReset_Mimex(TS ts)
272: {
273:   TS_Mimex       *mimex = (TS_Mimex*)ts->data;

277:   VecDestroy(&mimex->update);
278:   VecDestroy(&mimex->Xdot);
279:   return(0);
280: }

282: static PetscErrorCode TSDestroy_Mimex(TS ts)
283: {

287:   TSReset_Mimex(ts);
288:   PetscFree(ts->data);
289:   return(0);
290: }
291: /*------------------------------------------------------------*/

293: static PetscErrorCode TSSetFromOptions_Mimex(PetscOptionItems *PetscOptionsObject, TS ts)
294: {
295:   TS_Mimex      *mimex = (TS_Mimex *) ts->data;

299:   PetscOptionsHead(PetscOptionsObject, "MIMEX ODE solver options");
300:   {
301:     PetscOptionsInt("-ts_mimex_version", "Algorithm version", "TSMimexSetVersion", mimex->version, &mimex->version, NULL);
302:   }
303:   PetscOptionsTail();
304:   return(0);
305: }

307: static PetscErrorCode TSView_Mimex(TS ts,PetscViewer viewer)
308: {
309:   TS_Mimex      *mimex = (TS_Mimex *) ts->data;
310:   PetscBool      iascii;

314:   PetscObjectTypeCompare((PetscObject) viewer, PETSCVIEWERASCII, &iascii);
315:   if (iascii) {
316:     PetscViewerASCIIPrintf(viewer, "  Version = %D\n", mimex->version);
317:   }
318:   return(0);
319: }

321: static PetscErrorCode TSInterpolate_Mimex(TS ts,PetscReal t,Vec X)
322: {
323:   PetscReal      alpha = (ts->ptime - t)/ts->time_step;

327:   VecAXPBY(ts->vec_sol,1.0-alpha,alpha,X);
328:   return(0);
329: }

331: static PetscErrorCode TSComputeLinearStability_Mimex(TS ts,PetscReal xr,PetscReal xi,PetscReal *yr,PetscReal *yi)
332: {
334:   *yr = 1.0 + xr;
335:   *yi = xi;
336:   return(0);
337: }
338: /* ------------------------------------------------------------ */

340: /*MC
341:       TSMIMEX - ODE solver using the explicit forward Mimex method

343:   Level: beginner

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

347: M*/
348: PETSC_EXTERN PetscErrorCode TSCreate_Mimex(TS ts)
349: {
350:   TS_Mimex       *mimex;

354:   ts->ops->setup           = TSSetUp_Mimex;
355:   ts->ops->step            = TSStep_Mimex;
356:   ts->ops->reset           = TSReset_Mimex;
357:   ts->ops->destroy         = TSDestroy_Mimex;
358:   ts->ops->setfromoptions  = TSSetFromOptions_Mimex;
359:   ts->ops->view            = TSView_Mimex;
360:   ts->ops->interpolate     = TSInterpolate_Mimex;
361:   ts->ops->linearstability = TSComputeLinearStability_Mimex;
362:   ts->ops->snesfunction    = SNESTSFormFunction_Mimex;
363:   ts->ops->snesjacobian    = SNESTSFormJacobian_Mimex;
364:   ts->default_adapt_type   = TSADAPTNONE;

366:   PetscNewLog(ts,&mimex);
367:   ts->data = (void*)mimex;

369:   mimex->version = 1;
370:   return(0);
371: }