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, see `DMDATSRHSFunctionLocalFn` for the calling sequence
227: - ctx   - optional context for local residual evaluation

229:   Level: beginner

231: .seealso: [](ch_ts), `DMDA`, `DMDATSRHSFunctionLocalFn`, `TS`, `TSSetRHSFunction()`, `DMTSSetRHSFunction()`, `DMDATSSetRHSJacobianLocal()`, `DMDASNESSetFunctionLocal()`
232: @*/
233: PetscErrorCode DMDATSSetRHSFunctionLocal(DM dm, InsertMode imode, DMDATSRHSFunctionLocalFn *func, void *ctx)
234: {
235:   DMTS     sdm;
236:   DMTS_DA *dmdats;

238:   PetscFunctionBegin;
240:   PetscCall(DMGetDMTSWrite(dm, &sdm));
241:   PetscCall(DMDATSGetContext(dm, sdm, &dmdats));
242:   dmdats->rhsfunctionlocalimode = imode;
243:   dmdats->rhsfunctionlocal      = func;
244:   dmdats->rhsfunctionlocalctx   = ctx;
245:   PetscCall(DMTSSetRHSFunction(dm, TSComputeRHSFunction_DMDA, dmdats));
246:   PetscFunctionReturn(PETSC_SUCCESS);
247: }

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

252:   Logically Collective

254:   Input Parameters:
255: + dm   - `DM` to associate callback with
256: . func - local RHS Jacobian evaluation routine, see `DMDATSRHSJacobianLocalFn` for the calling sequence
257: - ctx  - optional context for local jacobian evaluation

259:   Level: beginner

261: .seealso: [](ch_ts), `DMDA`, `DMDATSRHSJacobianLocalFn`, `DMTSSetRHSJacobian()`,
262: `DMDATSSetRHSFunctionLocal()`, `DMDASNESSetJacobianLocal()`
263: @*/
264: PetscErrorCode DMDATSSetRHSJacobianLocal(DM dm, DMDATSRHSJacobianLocalFn *func, void *ctx)
265: {
266:   DMTS     sdm;
267:   DMTS_DA *dmdats;

269:   PetscFunctionBegin;
271:   PetscCall(DMGetDMTSWrite(dm, &sdm));
272:   PetscCall(DMDATSGetContext(dm, sdm, &dmdats));
273:   dmdats->rhsjacobianlocal    = func;
274:   dmdats->rhsjacobianlocalctx = ctx;
275:   PetscCall(DMTSSetRHSJacobian(dm, TSComputeRHSJacobian_DMDA, dmdats));
276:   PetscFunctionReturn(PETSC_SUCCESS);
277: }

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

282:   Logically Collective

284:   Input Parameters:
285: + dm    - `DM` to associate callback with
286: . imode - the insert mode of the function
287: . func  - local residual evaluation, see `DMDATSIFunctionLocalFn` for the calling sequence
288: - ctx   - optional context for local residual evaluation

290:   Level: beginner

292: .seealso: [](ch_ts), `DMDA`, `DMDATSIFunctionLocalFn`, `DMTSSetIFunction()`,
293: `DMDATSSetIJacobianLocal()`, `DMDASNESSetFunctionLocal()`
294: @*/
295: PetscErrorCode DMDATSSetIFunctionLocal(DM dm, InsertMode imode, DMDATSIFunctionLocalFn *func, void *ctx)
296: {
297:   DMTS     sdm;
298:   DMTS_DA *dmdats;

300:   PetscFunctionBegin;
302:   PetscCall(DMGetDMTSWrite(dm, &sdm));
303:   PetscCall(DMDATSGetContext(dm, sdm, &dmdats));
304:   dmdats->ifunctionlocalimode = imode;
305:   dmdats->ifunctionlocal      = func;
306:   dmdats->ifunctionlocalctx   = ctx;
307:   PetscCall(DMTSSetIFunction(dm, TSComputeIFunction_DMDA, dmdats));
308:   PetscFunctionReturn(PETSC_SUCCESS);
309: }

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

314:   Logically Collective

316:   Input Parameters:
317: + dm   - `DM` to associate callback with
318: . func - local residual evaluation, see `DMDATSIJacobianLocalFn` for the calling sequence
319: - ctx  - optional context for local residual evaluation

321:   Level: beginner

323: .seealso: [](ch_ts), `DMDA`, `DMDATSIJacobianLocalFn`, `DMTSSetJacobian()`,
324: `DMDATSSetIFunctionLocal()`, `DMDASNESSetJacobianLocal()`
325: @*/
326: PetscErrorCode DMDATSSetIJacobianLocal(DM dm, DMDATSIJacobianLocalFn *func, void *ctx)
327: {
328:   DMTS     sdm;
329:   DMTS_DA *dmdats;

331:   PetscFunctionBegin;
333:   PetscCall(DMGetDMTSWrite(dm, &sdm));
334:   PetscCall(DMDATSGetContext(dm, sdm, &dmdats));
335:   dmdats->ijacobianlocal    = func;
336:   dmdats->ijacobianlocalctx = ctx;
337:   PetscCall(DMTSSetIJacobian(dm, TSComputeIJacobian_DMDA, dmdats));
338:   PetscFunctionReturn(PETSC_SUCCESS);
339: }

341: PetscErrorCode TSMonitorDMDARayDestroy(void **mctx)
342: {
343:   TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx *)*mctx;

345:   PetscFunctionBegin;
346:   if (rayctx->lgctx) PetscCall(TSMonitorLGCtxDestroy(&rayctx->lgctx));
347:   PetscCall(VecDestroy(&rayctx->ray));
348:   PetscCall(VecScatterDestroy(&rayctx->scatter));
349:   PetscCall(PetscViewerDestroy(&rayctx->viewer));
350:   PetscCall(PetscFree(rayctx));
351:   PetscFunctionReturn(PETSC_SUCCESS);
352: }

354: PetscErrorCode TSMonitorDMDARay(TS ts, PetscInt steps, PetscReal time, Vec u, void *mctx)
355: {
356:   TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx *)mctx;
357:   Vec                  solution;

359:   PetscFunctionBegin;
360:   PetscCall(TSGetSolution(ts, &solution));
361:   PetscCall(VecScatterBegin(rayctx->scatter, solution, rayctx->ray, INSERT_VALUES, SCATTER_FORWARD));
362:   PetscCall(VecScatterEnd(rayctx->scatter, solution, rayctx->ray, INSERT_VALUES, SCATTER_FORWARD));
363:   if (rayctx->viewer) PetscCall(VecView(rayctx->ray, rayctx->viewer));
364:   PetscFunctionReturn(PETSC_SUCCESS);
365: }

367: PetscErrorCode TSMonitorLGDMDARay(TS ts, PetscInt step, PetscReal ptime, Vec u, void *ctx)
368: {
369:   TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx *)ctx;
370:   TSMonitorLGCtx       lgctx  = (TSMonitorLGCtx)rayctx->lgctx;
371:   Vec                  v      = rayctx->ray;
372:   const PetscScalar   *a;
373:   PetscInt             dim;

375:   PetscFunctionBegin;
376:   PetscCall(VecScatterBegin(rayctx->scatter, u, v, INSERT_VALUES, SCATTER_FORWARD));
377:   PetscCall(VecScatterEnd(rayctx->scatter, u, v, INSERT_VALUES, SCATTER_FORWARD));
378:   if (!step) {
379:     PetscDrawAxis axis;

381:     PetscCall(PetscDrawLGGetAxis(lgctx->lg, &axis));
382:     PetscCall(PetscDrawAxisSetLabels(axis, "Solution Ray as function of time", "Time", "Solution"));
383:     PetscCall(VecGetLocalSize(rayctx->ray, &dim));
384:     PetscCall(PetscDrawLGSetDimension(lgctx->lg, dim));
385:     PetscCall(PetscDrawLGReset(lgctx->lg));
386:   }
387:   PetscCall(VecGetArrayRead(v, &a));
388: #if defined(PETSC_USE_COMPLEX)
389:   {
390:     PetscReal *areal;
391:     PetscInt   i, n;
392:     PetscCall(VecGetLocalSize(v, &n));
393:     PetscCall(PetscMalloc1(n, &areal));
394:     for (i = 0; i < n; ++i) areal[i] = PetscRealPart(a[i]);
395:     PetscCall(PetscDrawLGAddCommonPoint(lgctx->lg, ptime, areal));
396:     PetscCall(PetscFree(areal));
397:   }
398: #else
399:   PetscCall(PetscDrawLGAddCommonPoint(lgctx->lg, ptime, a));
400: #endif
401:   PetscCall(VecRestoreArrayRead(v, &a));
402:   if (((lgctx->howoften > 0) && (!(step % lgctx->howoften))) || ((lgctx->howoften == -1) && ts->reason)) {
403:     PetscCall(PetscDrawLGDraw(lgctx->lg));
404:     PetscCall(PetscDrawLGSave(lgctx->lg));
405:   }
406:   PetscFunctionReturn(PETSC_SUCCESS);
407: }