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 of `func`:
230: $ PetscErrorCode func(DMDALocalInfo *info, PetscReal t, void *x, void *f, void *ctx)
231: +  info - defines the subdomain to evaluate the residual on
232: .  t - time at which to evaluate residual
233: .  x - array of local state information
234: .  f - output array of local residual information
235: -  ctx - optional user context

237:    Level: beginner

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

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

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

260:    Logically Collective

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

267:    Calling sequence of `func`:
268: $ PetscErrorCode func(DMDALocalInfo *info, PetscReal t, void* x, Mat J, Mat B, void *ctx)
269: +  info - defines the subdomain to evaluate the residual on
270: .  t    - time at which to evaluate residual
271: .  x    - array of local state information
272: .  J    - Jacobian matrix
273: .  B    - preconditioner matrix; often same as `J`
274: -  ctx  - optional context passed above

276:    Level: beginner

278: .seealso: [](ch_ts), `DMDA`, `DMTSSetRHSJacobian()`, `DMDATSSetRHSFunctionLocal()`, `DMDASNESSetJacobianLocal()`
279: @*/
280: PetscErrorCode DMDATSSetRHSJacobianLocal(DM dm, DMDATSRHSJacobianLocal func, void *ctx)
281: {
282:   DMTS     sdm;
283:   DMTS_DA *dmdats;

285:   PetscFunctionBegin;
287:   PetscCall(DMGetDMTSWrite(dm, &sdm));
288:   PetscCall(DMDATSGetContext(dm, sdm, &dmdats));
289:   dmdats->rhsjacobianlocal    = func;
290:   dmdats->rhsjacobianlocalctx = ctx;
291:   PetscCall(DMTSSetRHSJacobian(dm, TSComputeRHSJacobian_DMDA, dmdats));
292:   PetscFunctionReturn(PETSC_SUCCESS);
293: }

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

298:    Logically Collective

300:    Input Parameters:
301: +  dm   - `DM` to associate callback with
302: .  func - local residual evaluation
303: -  ctx  - optional context for local residual evaluation

305:    Calling sequence of `func`:
306: $  PetscErrorCode func(DMDALocalInfo *info, PetscReal t, Vec x, Vec xdot, Vec f, void *ctx)
307: +  info - defines the subdomain to evaluate the residual on
308: .  t    - time at which to evaluate residual
309: .  x    - array of local state information
310: .  xdot - array of local time derivative information
311: .  f    - output array of local function evaluation information
312: -  ctx - optional context passed above

314:    Level: beginner

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

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

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

337:    Logically Collective

339:    Input Parameters:
340: +  dm   - `DM` to associate callback with
341: .  func - local residual evaluation
342: -  ctx   - optional context for local residual evaluation

344:    Calling sequence of `func`:
345: $ PetscErrorCode func(DMDALocalInfo *info, PetscReal t, void* x, void *xdot, PetscScalar shift, Mat J, Mat B, void *ctx)
346: +  info - defines the subdomain to evaluate the residual on
347: .  t    - time at which to evaluate the jacobian
348: .  x    - array of local state information
349: .  xdot - time derivative at this state
350: .  shift - see `TSSetIJacobian()` for the meaning of this parameter
351: .  J    - Jacobian matrix
352: .  B    - preconditioner matrix; often same as `J`
353: -  ctx  - optional context passed above

355:    Level: beginner

357: .seealso: [](ch_ts), `DMDA`, `DMTSSetJacobian()`, `DMDATSSetIFunctionLocal()`, `DMDASNESSetJacobianLocal()`
358: @*/
359: PetscErrorCode DMDATSSetIJacobianLocal(DM dm, DMDATSIJacobianLocal func, void *ctx)
360: {
361:   DMTS     sdm;
362:   DMTS_DA *dmdats;

364:   PetscFunctionBegin;
366:   PetscCall(DMGetDMTSWrite(dm, &sdm));
367:   PetscCall(DMDATSGetContext(dm, sdm, &dmdats));
368:   dmdats->ijacobianlocal    = func;
369:   dmdats->ijacobianlocalctx = ctx;
370:   PetscCall(DMTSSetIJacobian(dm, TSComputeIJacobian_DMDA, dmdats));
371:   PetscFunctionReturn(PETSC_SUCCESS);
372: }

374: PetscErrorCode TSMonitorDMDARayDestroy(void **mctx)
375: {
376:   TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx *)*mctx;

378:   PetscFunctionBegin;
379:   if (rayctx->lgctx) PetscCall(TSMonitorLGCtxDestroy(&rayctx->lgctx));
380:   PetscCall(VecDestroy(&rayctx->ray));
381:   PetscCall(VecScatterDestroy(&rayctx->scatter));
382:   PetscCall(PetscViewerDestroy(&rayctx->viewer));
383:   PetscCall(PetscFree(rayctx));
384:   PetscFunctionReturn(PETSC_SUCCESS);
385: }

387: PetscErrorCode TSMonitorDMDARay(TS ts, PetscInt steps, PetscReal time, Vec u, void *mctx)
388: {
389:   TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx *)mctx;
390:   Vec                  solution;

392:   PetscFunctionBegin;
393:   PetscCall(TSGetSolution(ts, &solution));
394:   PetscCall(VecScatterBegin(rayctx->scatter, solution, rayctx->ray, INSERT_VALUES, SCATTER_FORWARD));
395:   PetscCall(VecScatterEnd(rayctx->scatter, solution, rayctx->ray, INSERT_VALUES, SCATTER_FORWARD));
396:   if (rayctx->viewer) PetscCall(VecView(rayctx->ray, rayctx->viewer));
397:   PetscFunctionReturn(PETSC_SUCCESS);
398: }

400: PetscErrorCode TSMonitorLGDMDARay(TS ts, PetscInt step, PetscReal ptime, Vec u, void *ctx)
401: {
402:   TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx *)ctx;
403:   TSMonitorLGCtx       lgctx  = (TSMonitorLGCtx)rayctx->lgctx;
404:   Vec                  v      = rayctx->ray;
405:   const PetscScalar   *a;
406:   PetscInt             dim;

408:   PetscFunctionBegin;
409:   PetscCall(VecScatterBegin(rayctx->scatter, u, v, INSERT_VALUES, SCATTER_FORWARD));
410:   PetscCall(VecScatterEnd(rayctx->scatter, u, v, INSERT_VALUES, SCATTER_FORWARD));
411:   if (!step) {
412:     PetscDrawAxis axis;

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