Actual source code: dmdats.c

  1: #include <petscdmda.h>
  2: #include <petsc/private/dmimpl.h>
  3: #include <petsc/private/tsimpl.h>
  4: #include <petscdraw.h>

  6: /* This structure holds the user-provided DMDA callbacks */
  7: typedef struct {
  8:   PetscErrorCode (*ifunctionlocal)(DMDALocalInfo *, PetscReal, void *, void *, void *, void *);
  9:   PetscErrorCode (*rhsfunctionlocal)(DMDALocalInfo *, PetscReal, void *, void *, void *);
 10:   PetscErrorCode (*ijacobianlocal)(DMDALocalInfo *, PetscReal, void *, void *, PetscReal, Mat, Mat, void *);
 11:   PetscErrorCode (*rhsjacobianlocal)(DMDALocalInfo *, PetscReal, void *, Mat, Mat, void *);
 12:   void      *ifunctionlocalctx;
 13:   void      *ijacobianlocalctx;
 14:   void      *rhsfunctionlocalctx;
 15:   void      *rhsjacobianlocalctx;
 16:   InsertMode ifunctionlocalimode;
 17:   InsertMode rhsfunctionlocalimode;
 18: } DMTS_DA;

 20: static PetscErrorCode DMTSDestroy_DMDA(DMTS sdm)
 21: {
 22:   PetscFree(sdm->data);
 23:   return 0;
 24: }

 26: static PetscErrorCode DMTSDuplicate_DMDA(DMTS oldsdm, DMTS sdm)
 27: {
 28:   PetscNew((DMTS_DA **)&sdm->data);
 29:   if (oldsdm->data) PetscMemcpy(sdm->data, oldsdm->data, sizeof(DMTS_DA));
 30:   return 0;
 31: }

 33: static PetscErrorCode DMDATSGetContext(DM dm, DMTS sdm, DMTS_DA **dmdats)
 34: {
 35:   *dmdats = NULL;
 36:   if (!sdm->data) {
 37:     PetscNew((DMTS_DA **)&sdm->data);
 38:     sdm->ops->destroy   = DMTSDestroy_DMDA;
 39:     sdm->ops->duplicate = DMTSDuplicate_DMDA;
 40:   }
 41:   *dmdats = (DMTS_DA *)sdm->data;
 42:   return 0;
 43: }

 45: static PetscErrorCode TSComputeIFunction_DMDA(TS ts, PetscReal ptime, Vec X, Vec Xdot, Vec F, void *ctx)
 46: {
 47:   DM            dm;
 48:   DMTS_DA      *dmdats = (DMTS_DA *)ctx;
 49:   DMDALocalInfo info;
 50:   Vec           Xloc, Xdotloc;
 51:   void         *x, *f, *xdot;

 57:   TSGetDM(ts, &dm);
 58:   DMGetLocalVector(dm, &Xdotloc);
 59:   DMGlobalToLocalBegin(dm, Xdot, INSERT_VALUES, Xdotloc);
 60:   DMGlobalToLocalEnd(dm, Xdot, INSERT_VALUES, Xdotloc);
 61:   DMGetLocalVector(dm, &Xloc);
 62:   DMGlobalToLocalBegin(dm, X, INSERT_VALUES, Xloc);
 63:   DMGlobalToLocalEnd(dm, X, INSERT_VALUES, Xloc);
 64:   DMDAGetLocalInfo(dm, &info);
 65:   DMDAVecGetArray(dm, Xloc, &x);
 66:   DMDAVecGetArray(dm, Xdotloc, &xdot);
 67:   switch (dmdats->ifunctionlocalimode) {
 68:   case INSERT_VALUES: {
 69:     DMDAVecGetArray(dm, F, &f);
 70:     CHKMEMQ;
 71:     (*dmdats->ifunctionlocal)(&info, ptime, x, xdot, f, dmdats->ifunctionlocalctx);
 72:     CHKMEMQ;
 73:     DMDAVecRestoreArray(dm, F, &f);
 74:   } break;
 75:   case ADD_VALUES: {
 76:     Vec Floc;
 77:     DMGetLocalVector(dm, &Floc);
 78:     VecZeroEntries(Floc);
 79:     DMDAVecGetArray(dm, Floc, &f);
 80:     CHKMEMQ;
 81:     (*dmdats->ifunctionlocal)(&info, ptime, x, xdot, f, dmdats->ifunctionlocalctx);
 82:     CHKMEMQ;
 83:     DMDAVecRestoreArray(dm, Floc, &f);
 84:     VecZeroEntries(F);
 85:     DMLocalToGlobalBegin(dm, Floc, ADD_VALUES, F);
 86:     DMLocalToGlobalEnd(dm, Floc, ADD_VALUES, F);
 87:     DMRestoreLocalVector(dm, &Floc);
 88:   } break;
 89:   default:
 90:     SETERRQ(PetscObjectComm((PetscObject)ts), PETSC_ERR_ARG_INCOMP, "Cannot use imode=%d", (int)dmdats->ifunctionlocalimode);
 91:   }
 92:   DMDAVecRestoreArray(dm, Xloc, &x);
 93:   DMRestoreLocalVector(dm, &Xloc);
 94:   DMDAVecRestoreArray(dm, Xdotloc, &xdot);
 95:   DMRestoreLocalVector(dm, &Xdotloc);
 96:   return 0;
 97: }

 99: static PetscErrorCode TSComputeIJacobian_DMDA(TS ts, PetscReal ptime, Vec X, Vec Xdot, PetscReal shift, Mat A, Mat B, void *ctx)
100: {
101:   DM            dm;
102:   DMTS_DA      *dmdats = (DMTS_DA *)ctx;
103:   DMDALocalInfo info;
104:   Vec           Xloc;
105:   void         *x, *xdot;

108:   TSGetDM(ts, &dm);

110:   if (dmdats->ijacobianlocal) {
111:     DMGetLocalVector(dm, &Xloc);
112:     DMGlobalToLocalBegin(dm, X, INSERT_VALUES, Xloc);
113:     DMGlobalToLocalEnd(dm, X, INSERT_VALUES, Xloc);
114:     DMDAGetLocalInfo(dm, &info);
115:     DMDAVecGetArray(dm, Xloc, &x);
116:     DMDAVecGetArray(dm, Xdot, &xdot);
117:     CHKMEMQ;
118:     (*dmdats->ijacobianlocal)(&info, ptime, x, xdot, shift, A, B, dmdats->ijacobianlocalctx);
119:     CHKMEMQ;
120:     DMDAVecRestoreArray(dm, Xloc, &x);
121:     DMDAVecRestoreArray(dm, Xdot, &xdot);
122:     DMRestoreLocalVector(dm, &Xloc);
123:   } else SETERRQ(PetscObjectComm((PetscObject)ts), PETSC_ERR_PLIB, "TSComputeIJacobian_DMDA() called without calling DMDATSSetIJacobian()");
124:   /* This will be redundant if the user called both, but it's too common to forget. */
125:   if (A != B) {
126:     MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY);
127:     MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY);
128:   }
129:   return 0;
130: }

132: static PetscErrorCode TSComputeRHSFunction_DMDA(TS ts, PetscReal ptime, Vec X, Vec F, void *ctx)
133: {
134:   DM            dm;
135:   DMTS_DA      *dmdats = (DMTS_DA *)ctx;
136:   DMDALocalInfo info;
137:   Vec           Xloc;
138:   void         *x, *f;

144:   TSGetDM(ts, &dm);
145:   DMGetLocalVector(dm, &Xloc);
146:   DMGlobalToLocalBegin(dm, X, INSERT_VALUES, Xloc);
147:   DMGlobalToLocalEnd(dm, X, INSERT_VALUES, Xloc);
148:   DMDAGetLocalInfo(dm, &info);
149:   DMDAVecGetArray(dm, Xloc, &x);
150:   switch (dmdats->rhsfunctionlocalimode) {
151:   case INSERT_VALUES: {
152:     DMDAVecGetArray(dm, F, &f);
153:     CHKMEMQ;
154:     (*dmdats->rhsfunctionlocal)(&info, ptime, x, f, dmdats->rhsfunctionlocalctx);
155:     CHKMEMQ;
156:     DMDAVecRestoreArray(dm, F, &f);
157:   } break;
158:   case ADD_VALUES: {
159:     Vec Floc;
160:     DMGetLocalVector(dm, &Floc);
161:     VecZeroEntries(Floc);
162:     DMDAVecGetArray(dm, Floc, &f);
163:     CHKMEMQ;
164:     (*dmdats->rhsfunctionlocal)(&info, ptime, x, f, dmdats->rhsfunctionlocalctx);
165:     CHKMEMQ;
166:     DMDAVecRestoreArray(dm, Floc, &f);
167:     VecZeroEntries(F);
168:     DMLocalToGlobalBegin(dm, Floc, ADD_VALUES, F);
169:     DMLocalToGlobalEnd(dm, Floc, ADD_VALUES, F);
170:     DMRestoreLocalVector(dm, &Floc);
171:   } break;
172:   default:
173:     SETERRQ(PetscObjectComm((PetscObject)ts), PETSC_ERR_ARG_INCOMP, "Cannot use imode=%d", (int)dmdats->rhsfunctionlocalimode);
174:   }
175:   DMDAVecRestoreArray(dm, Xloc, &x);
176:   DMRestoreLocalVector(dm, &Xloc);
177:   return 0;
178: }

180: static PetscErrorCode TSComputeRHSJacobian_DMDA(TS ts, PetscReal ptime, Vec X, Mat A, Mat B, void *ctx)
181: {
182:   DM            dm;
183:   DMTS_DA      *dmdats = (DMTS_DA *)ctx;
184:   DMDALocalInfo info;
185:   Vec           Xloc;
186:   void         *x;

189:   TSGetDM(ts, &dm);

191:   if (dmdats->rhsjacobianlocal) {
192:     DMGetLocalVector(dm, &Xloc);
193:     DMGlobalToLocalBegin(dm, X, INSERT_VALUES, Xloc);
194:     DMGlobalToLocalEnd(dm, X, INSERT_VALUES, Xloc);
195:     DMDAGetLocalInfo(dm, &info);
196:     DMDAVecGetArray(dm, Xloc, &x);
197:     CHKMEMQ;
198:     (*dmdats->rhsjacobianlocal)(&info, ptime, x, A, B, dmdats->rhsjacobianlocalctx);
199:     CHKMEMQ;
200:     DMDAVecRestoreArray(dm, Xloc, &x);
201:     DMRestoreLocalVector(dm, &Xloc);
202:   } else SETERRQ(PetscObjectComm((PetscObject)ts), PETSC_ERR_PLIB, "TSComputeRHSJacobian_DMDA() called without calling DMDATSSetRHSJacobian()");
203:   /* This will be redundant if the user called both, but it's too common to forget. */
204:   if (A != B) {
205:     MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY);
206:     MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY);
207:   }
208:   return 0;
209: }

211: /*@C
212:    DMDATSSetRHSFunctionLocal - set a local residual evaluation function for use with `DMDA`

214:    Logically Collective

216:    Input Parameters:
217: +  dm - `DM` to associate callback with
218: .  imode - insert mode for the residual
219: .  func - local residual evaluation
220: -  ctx - optional context for local residual evaluation

222:    Calling sequence for func:

224: $ func(DMDALocalInfo info,PetscReal t,void *x,void *f,void *ctx)

226: +  info - `DMDALocalInfo` defining the subdomain to evaluate the residual on
227: .  t - time at which to evaluate residual
228: .  x - array of local state information
229: .  f - output array of local residual information
230: -  ctx - optional user context

232:    Level: beginner

234: .seealso: [](chapter_ts), `DMDA`, `TS`, `TSSetRHSFunction()`, `DMTSSetRHSFunction()`, `DMDATSSetRHSJacobianLocal()`, `DMDASNESSetFunctionLocal()`
235: @*/
236: PetscErrorCode DMDATSSetRHSFunctionLocal(DM dm, InsertMode imode, DMDATSRHSFunctionLocal func, void *ctx)
237: {
238:   DMTS     sdm;
239:   DMTS_DA *dmdats;

242:   DMGetDMTSWrite(dm, &sdm);
243:   DMDATSGetContext(dm, sdm, &dmdats);
244:   dmdats->rhsfunctionlocalimode = imode;
245:   dmdats->rhsfunctionlocal      = func;
246:   dmdats->rhsfunctionlocalctx   = ctx;
247:   DMTSSetRHSFunction(dm, TSComputeRHSFunction_DMDA, dmdats);
248:   return 0;
249: }

251: /*@C
252:    DMDATSSetRHSJacobianLocal - set a local residual evaluation function for use with `DMDA`

254:    Logically Collective

256:    Input Parameters:
257: +  dm    - `DM` to associate callback with
258: .  func  - local RHS Jacobian evaluation routine
259: -  ctx   - optional context for local jacobian evaluation

261:    Calling sequence for func:

263: $ func(DMDALocalInfo* info,PetscReal t,void* x,Mat J,Mat B,void *ctx);

265: +  info - `DMDALocalInfo` defining the subdomain to evaluate the residual on
266: .  t    - time at which to evaluate residual
267: .  x    - array of local state information
268: .  J    - Jacobian matrix
269: .  B    - preconditioner matrix; often same as J
270: -  ctx  - optional context passed above

272:    Level: beginner

274: .seealso: [](chapter_ts), `DMDA`, `DMTSSetRHSJacobian()`, `DMDATSSetRHSFunctionLocal()`, `DMDASNESSetJacobianLocal()`
275: @*/
276: PetscErrorCode DMDATSSetRHSJacobianLocal(DM dm, DMDATSRHSJacobianLocal func, void *ctx)
277: {
278:   DMTS     sdm;
279:   DMTS_DA *dmdats;

282:   DMGetDMTSWrite(dm, &sdm);
283:   DMDATSGetContext(dm, sdm, &dmdats);
284:   dmdats->rhsjacobianlocal    = func;
285:   dmdats->rhsjacobianlocalctx = ctx;
286:   DMTSSetRHSJacobian(dm, TSComputeRHSJacobian_DMDA, dmdats);
287:   return 0;
288: }

290: /*@C
291:    DMDATSSetIFunctionLocal - set a local residual evaluation function for use with `DMDA`

293:    Logically Collective

295:    Input Parameters:
296: +  dm   - `DM` to associate callback with
297: .  func - local residual evaluation
298: -  ctx  - optional context for local residual evaluation

300:    Calling sequence for func:
301: +  info - `DMDALocalInfo` defining the subdomain to evaluate the residual on
302: .  t    - time at which to evaluate residual
303: .  x    - array of local state information
304: .  xdot - array of local time derivative information
305: .  f    - output array of local function evaluation information
306: -  ctx - optional context passed above

308:    Level: beginner

310: .seealso: [](chapter_ts), `DMDA`, `DMTSSetIFunction()`, `DMDATSSetIJacobianLocal()`, `DMDASNESSetFunctionLocal()`
311: @*/
312: PetscErrorCode DMDATSSetIFunctionLocal(DM dm, InsertMode imode, DMDATSIFunctionLocal func, void *ctx)
313: {
314:   DMTS     sdm;
315:   DMTS_DA *dmdats;

318:   DMGetDMTSWrite(dm, &sdm);
319:   DMDATSGetContext(dm, sdm, &dmdats);
320:   dmdats->ifunctionlocalimode = imode;
321:   dmdats->ifunctionlocal      = func;
322:   dmdats->ifunctionlocalctx   = ctx;
323:   DMTSSetIFunction(dm, TSComputeIFunction_DMDA, dmdats);
324:   return 0;
325: }

327: /*@C
328:    DMDATSSetIJacobianLocal - set a local residual evaluation function for use with `DMDA`

330:    Logically Collective

332:    Input Parameters:
333: +  dm   - `DM` to associate callback with
334: .  func - local residual evaluation
335: -  ctx   - optional context for local residual evaluation

337:    Calling sequence for func:

339: $ func(DMDALocalInfo* info,PetscReal t,void* x,void *xdot,PetscScalar shift,Mat J,Mat B,void *ctx);

341: +  info - `DMDALocalInfo` defining the subdomain to evaluate the residual on
342: .  t    - time at which to evaluate the jacobian
343: .  x    - array of local state information
344: .  xdot - time derivative at this state
345: .  shift - see TSSetIJacobian() for the meaning of this parameter
346: .  J    - Jacobian matrix
347: .  B    - preconditioner matrix; often same as J
348: -  ctx  - optional context passed above

350:    Level: beginner

352: .seealso: [](chapter_ts), `DMDA`, `DMTSSetJacobian()`, `DMDATSSetIFunctionLocal()`, `DMDASNESSetJacobianLocal()`
353: @*/
354: PetscErrorCode DMDATSSetIJacobianLocal(DM dm, DMDATSIJacobianLocal func, void *ctx)
355: {
356:   DMTS     sdm;
357:   DMTS_DA *dmdats;

360:   DMGetDMTSWrite(dm, &sdm);
361:   DMDATSGetContext(dm, sdm, &dmdats);
362:   dmdats->ijacobianlocal    = func;
363:   dmdats->ijacobianlocalctx = ctx;
364:   DMTSSetIJacobian(dm, TSComputeIJacobian_DMDA, dmdats);
365:   return 0;
366: }

368: PetscErrorCode TSMonitorDMDARayDestroy(void **mctx)
369: {
370:   TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx *)*mctx;

372:   if (rayctx->lgctx) TSMonitorLGCtxDestroy(&rayctx->lgctx);
373:   VecDestroy(&rayctx->ray);
374:   VecScatterDestroy(&rayctx->scatter);
375:   PetscViewerDestroy(&rayctx->viewer);
376:   PetscFree(rayctx);
377:   return 0;
378: }

380: PetscErrorCode TSMonitorDMDARay(TS ts, PetscInt steps, PetscReal time, Vec u, void *mctx)
381: {
382:   TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx *)mctx;
383:   Vec                  solution;

385:   TSGetSolution(ts, &solution);
386:   VecScatterBegin(rayctx->scatter, solution, rayctx->ray, INSERT_VALUES, SCATTER_FORWARD);
387:   VecScatterEnd(rayctx->scatter, solution, rayctx->ray, INSERT_VALUES, SCATTER_FORWARD);
388:   if (rayctx->viewer) VecView(rayctx->ray, rayctx->viewer);
389:   return 0;
390: }

392: PetscErrorCode TSMonitorLGDMDARay(TS ts, PetscInt step, PetscReal ptime, Vec u, void *ctx)
393: {
394:   TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx *)ctx;
395:   TSMonitorLGCtx       lgctx  = (TSMonitorLGCtx)rayctx->lgctx;
396:   Vec                  v      = rayctx->ray;
397:   const PetscScalar   *a;
398:   PetscInt             dim;

400:   VecScatterBegin(rayctx->scatter, u, v, INSERT_VALUES, SCATTER_FORWARD);
401:   VecScatterEnd(rayctx->scatter, u, v, INSERT_VALUES, SCATTER_FORWARD);
402:   if (!step) {
403:     PetscDrawAxis axis;

405:     PetscDrawLGGetAxis(lgctx->lg, &axis);
406:     PetscDrawAxisSetLabels(axis, "Solution Ray as function of time", "Time", "Solution");
407:     VecGetLocalSize(rayctx->ray, &dim);
408:     PetscDrawLGSetDimension(lgctx->lg, dim);
409:     PetscDrawLGReset(lgctx->lg);
410:   }
411:   VecGetArrayRead(v, &a);
412: #if defined(PETSC_USE_COMPLEX)
413:   {
414:     PetscReal *areal;
415:     PetscInt   i, n;
416:     VecGetLocalSize(v, &n);
417:     PetscMalloc1(n, &areal);
418:     for (i = 0; i < n; ++i) areal[i] = PetscRealPart(a[i]);
419:     PetscDrawLGAddCommonPoint(lgctx->lg, ptime, areal);
420:     PetscFree(areal);
421:   }
422: #else
423:   PetscDrawLGAddCommonPoint(lgctx->lg, ptime, a);
424: #endif
425:   VecRestoreArrayRead(v, &a);
426:   if (((lgctx->howoften > 0) && (!(step % lgctx->howoften))) || ((lgctx->howoften == -1) && ts->reason)) {
427:     PetscDrawLGDraw(lgctx->lg);
428:     PetscDrawLGSave(lgctx->lg);
429:   }
430:   return 0;
431: }