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:   PetscFunctionBegin;
 23:   PetscCall(PetscFree(sdm->data));
 24:   PetscFunctionReturn(PETSC_SUCCESS);
 25: }

 27: static PetscErrorCode DMTSDuplicate_DMDA(DMTS oldsdm, DMTS sdm)
 28: {
 29:   PetscFunctionBegin;
 30:   PetscCall(PetscNew((DMTS_DA **)&sdm->data));
 31:   if (oldsdm->data) PetscCall(PetscMemcpy(sdm->data, oldsdm->data, sizeof(DMTS_DA)));
 32:   PetscFunctionReturn(PETSC_SUCCESS);
 33: }

 35: static PetscErrorCode DMDATSGetContext(DM dm, DMTS sdm, DMTS_DA **dmdats)
 36: {
 37:   PetscFunctionBegin;
 38:   *dmdats = NULL;
 39:   if (!sdm->data) {
 40:     PetscCall(PetscNew((DMTS_DA **)&sdm->data));
 41:     sdm->ops->destroy   = DMTSDestroy_DMDA;
 42:     sdm->ops->duplicate = DMTSDuplicate_DMDA;
 43:   }
 44:   *dmdats = (DMTS_DA *)sdm->data;
 45:   PetscFunctionReturn(PETSC_SUCCESS);
 46: }

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

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

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

111:   PetscFunctionBegin;
112:   PetscCheck(dmdats->ifunctionlocal, PetscObjectComm((PetscObject)ts), PETSC_ERR_PLIB, "Corrupt context");
113:   PetscCall(TSGetDM(ts, &dm));

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

137: static PetscErrorCode TSComputeRHSFunction_DMDA(TS ts, PetscReal ptime, Vec X, Vec F, void *ctx)
138: {
139:   DM            dm;
140:   DMTS_DA      *dmdats = (DMTS_DA *)ctx;
141:   DMDALocalInfo info;
142:   Vec           Xloc;
143:   void         *x, *f;

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

186: static PetscErrorCode TSComputeRHSJacobian_DMDA(TS ts, PetscReal ptime, Vec X, Mat A, Mat B, void *ctx)
187: {
188:   DM            dm;
189:   DMTS_DA      *dmdats = (DMTS_DA *)ctx;
190:   DMDALocalInfo info;
191:   Vec           Xloc;
192:   void         *x;

194:   PetscFunctionBegin;
195:   PetscCheck(dmdats->rhsfunctionlocal, PetscObjectComm((PetscObject)ts), PETSC_ERR_PLIB, "Corrupt context");
196:   PetscCall(TSGetDM(ts, &dm));

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

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

221:    Logically Collective

223:    Input Parameters:
224: +  dm - `DM` to associate callback with
225: .  imode - insert mode for the residual
226: .  func - local residual evaluation
227: -  ctx - optional context for local residual evaluation

229:    Calling sequence for func:

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

233: +  info - `DMDALocalInfo` defining the subdomain to evaluate the residual on
234: .  t - time at which to evaluate residual
235: .  x - array of local state information
236: .  f - output array of local residual information
237: -  ctx - optional user context

239:    Level: beginner

241: .seealso: [](chapter_ts), `DMDA`, `TS`, `TSSetRHSFunction()`, `DMTSSetRHSFunction()`, `DMDATSSetRHSJacobianLocal()`, `DMDASNESSetFunctionLocal()`
242: @*/
243: PetscErrorCode DMDATSSetRHSFunctionLocal(DM dm, InsertMode imode, DMDATSRHSFunctionLocal func, void *ctx)
244: {
245:   DMTS     sdm;
246:   DMTS_DA *dmdats;

248:   PetscFunctionBegin;
250:   PetscCall(DMGetDMTSWrite(dm, &sdm));
251:   PetscCall(DMDATSGetContext(dm, sdm, &dmdats));
252:   dmdats->rhsfunctionlocalimode = imode;
253:   dmdats->rhsfunctionlocal      = func;
254:   dmdats->rhsfunctionlocalctx   = ctx;
255:   PetscCall(DMTSSetRHSFunction(dm, TSComputeRHSFunction_DMDA, dmdats));
256:   PetscFunctionReturn(PETSC_SUCCESS);
257: }

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

262:    Logically Collective

264:    Input Parameters:
265: +  dm    - `DM` to associate callback with
266: .  func  - local RHS Jacobian evaluation routine
267: -  ctx   - optional context for local jacobian evaluation

269:    Calling sequence for func:

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

273: +  info - `DMDALocalInfo` defining the subdomain to evaluate the residual on
274: .  t    - time at which to evaluate residual
275: .  x    - array of local state information
276: .  J    - Jacobian matrix
277: .  B    - preconditioner matrix; often same as J
278: -  ctx  - optional context passed above

280:    Level: beginner

282: .seealso: [](chapter_ts), `DMDA`, `DMTSSetRHSJacobian()`, `DMDATSSetRHSFunctionLocal()`, `DMDASNESSetJacobianLocal()`
283: @*/
284: PetscErrorCode DMDATSSetRHSJacobianLocal(DM dm, DMDATSRHSJacobianLocal func, void *ctx)
285: {
286:   DMTS     sdm;
287:   DMTS_DA *dmdats;

289:   PetscFunctionBegin;
291:   PetscCall(DMGetDMTSWrite(dm, &sdm));
292:   PetscCall(DMDATSGetContext(dm, sdm, &dmdats));
293:   dmdats->rhsjacobianlocal    = func;
294:   dmdats->rhsjacobianlocalctx = ctx;
295:   PetscCall(DMTSSetRHSJacobian(dm, TSComputeRHSJacobian_DMDA, dmdats));
296:   PetscFunctionReturn(PETSC_SUCCESS);
297: }

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

302:    Logically Collective

304:    Input Parameters:
305: +  dm   - `DM` to associate callback with
306: .  func - local residual evaluation
307: -  ctx  - optional context for local residual evaluation

309:    Calling sequence for func:
310: +  info - `DMDALocalInfo` defining the subdomain to evaluate the residual on
311: .  t    - time at which to evaluate residual
312: .  x    - array of local state information
313: .  xdot - array of local time derivative information
314: .  f    - output array of local function evaluation information
315: -  ctx - optional context passed above

317:    Level: beginner

319: .seealso: [](chapter_ts), `DMDA`, `DMTSSetIFunction()`, `DMDATSSetIJacobianLocal()`, `DMDASNESSetFunctionLocal()`
320: @*/
321: PetscErrorCode DMDATSSetIFunctionLocal(DM dm, InsertMode imode, DMDATSIFunctionLocal func, void *ctx)
322: {
323:   DMTS     sdm;
324:   DMTS_DA *dmdats;

326:   PetscFunctionBegin;
328:   PetscCall(DMGetDMTSWrite(dm, &sdm));
329:   PetscCall(DMDATSGetContext(dm, sdm, &dmdats));
330:   dmdats->ifunctionlocalimode = imode;
331:   dmdats->ifunctionlocal      = func;
332:   dmdats->ifunctionlocalctx   = ctx;
333:   PetscCall(DMTSSetIFunction(dm, TSComputeIFunction_DMDA, dmdats));
334:   PetscFunctionReturn(PETSC_SUCCESS);
335: }

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

340:    Logically Collective

342:    Input Parameters:
343: +  dm   - `DM` to associate callback with
344: .  func - local residual evaluation
345: -  ctx   - optional context for local residual evaluation

347:    Calling sequence for func:

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

351: +  info - `DMDALocalInfo` defining the subdomain to evaluate the residual on
352: .  t    - time at which to evaluate the jacobian
353: .  x    - array of local state information
354: .  xdot - time derivative at this state
355: .  shift - see TSSetIJacobian() for the meaning of this parameter
356: .  J    - Jacobian matrix
357: .  B    - preconditioner matrix; often same as J
358: -  ctx  - optional context passed above

360:    Level: beginner

362: .seealso: [](chapter_ts), `DMDA`, `DMTSSetJacobian()`, `DMDATSSetIFunctionLocal()`, `DMDASNESSetJacobianLocal()`
363: @*/
364: PetscErrorCode DMDATSSetIJacobianLocal(DM dm, DMDATSIJacobianLocal func, void *ctx)
365: {
366:   DMTS     sdm;
367:   DMTS_DA *dmdats;

369:   PetscFunctionBegin;
371:   PetscCall(DMGetDMTSWrite(dm, &sdm));
372:   PetscCall(DMDATSGetContext(dm, sdm, &dmdats));
373:   dmdats->ijacobianlocal    = func;
374:   dmdats->ijacobianlocalctx = ctx;
375:   PetscCall(DMTSSetIJacobian(dm, TSComputeIJacobian_DMDA, dmdats));
376:   PetscFunctionReturn(PETSC_SUCCESS);
377: }

379: PetscErrorCode TSMonitorDMDARayDestroy(void **mctx)
380: {
381:   TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx *)*mctx;

383:   PetscFunctionBegin;
384:   if (rayctx->lgctx) PetscCall(TSMonitorLGCtxDestroy(&rayctx->lgctx));
385:   PetscCall(VecDestroy(&rayctx->ray));
386:   PetscCall(VecScatterDestroy(&rayctx->scatter));
387:   PetscCall(PetscViewerDestroy(&rayctx->viewer));
388:   PetscCall(PetscFree(rayctx));
389:   PetscFunctionReturn(PETSC_SUCCESS);
390: }

392: PetscErrorCode TSMonitorDMDARay(TS ts, PetscInt steps, PetscReal time, Vec u, void *mctx)
393: {
394:   TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx *)mctx;
395:   Vec                  solution;

397:   PetscFunctionBegin;
398:   PetscCall(TSGetSolution(ts, &solution));
399:   PetscCall(VecScatterBegin(rayctx->scatter, solution, rayctx->ray, INSERT_VALUES, SCATTER_FORWARD));
400:   PetscCall(VecScatterEnd(rayctx->scatter, solution, rayctx->ray, INSERT_VALUES, SCATTER_FORWARD));
401:   if (rayctx->viewer) PetscCall(VecView(rayctx->ray, rayctx->viewer));
402:   PetscFunctionReturn(PETSC_SUCCESS);
403: }

405: PetscErrorCode TSMonitorLGDMDARay(TS ts, PetscInt step, PetscReal ptime, Vec u, void *ctx)
406: {
407:   TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx *)ctx;
408:   TSMonitorLGCtx       lgctx  = (TSMonitorLGCtx)rayctx->lgctx;
409:   Vec                  v      = rayctx->ray;
410:   const PetscScalar   *a;
411:   PetscInt             dim;

413:   PetscFunctionBegin;
414:   PetscCall(VecScatterBegin(rayctx->scatter, u, v, INSERT_VALUES, SCATTER_FORWARD));
415:   PetscCall(VecScatterEnd(rayctx->scatter, u, v, INSERT_VALUES, SCATTER_FORWARD));
416:   if (!step) {
417:     PetscDrawAxis axis;

419:     PetscCall(PetscDrawLGGetAxis(lgctx->lg, &axis));
420:     PetscCall(PetscDrawAxisSetLabels(axis, "Solution Ray as function of time", "Time", "Solution"));
421:     PetscCall(VecGetLocalSize(rayctx->ray, &dim));
422:     PetscCall(PetscDrawLGSetDimension(lgctx->lg, dim));
423:     PetscCall(PetscDrawLGReset(lgctx->lg));
424:   }
425:   PetscCall(VecGetArrayRead(v, &a));
426: #if defined(PETSC_USE_COMPLEX)
427:   {
428:     PetscReal *areal;
429:     PetscInt   i, n;
430:     PetscCall(VecGetLocalSize(v, &n));
431:     PetscCall(PetscMalloc1(n, &areal));
432:     for (i = 0; i < n; ++i) areal[i] = PetscRealPart(a[i]);
433:     PetscCall(PetscDrawLGAddCommonPoint(lgctx->lg, ptime, areal));
434:     PetscCall(PetscFree(areal));
435:   }
436: #else
437:   PetscCall(PetscDrawLGAddCommonPoint(lgctx->lg, ptime, a));
438: #endif
439:   PetscCall(VecRestoreArrayRead(v, &a));
440:   if (((lgctx->howoften > 0) && (!(step % lgctx->howoften))) || ((lgctx->howoften == -1) && ts->reason)) {
441:     PetscCall(PetscDrawLGDraw(lgctx->lg));
442:     PetscCall(PetscDrawLGSave(lgctx->lg));
443:   }
444:   PetscFunctionReturn(PETSC_SUCCESS);
445: }