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, Xdotloc;
109: void *x, *xdot;
111: PetscFunctionBegin;
112: PetscCheck(dmdats->ifunctionlocal, PetscObjectComm((PetscObject)ts), PETSC_ERR_PLIB, "Corrupt context");
113: PetscCall(TSGetDM(ts, &dm));
115: PetscCheck(dmdats->ijacobianlocal, PetscObjectComm((PetscObject)ts), PETSC_ERR_PLIB, "TSComputeIJacobian_DMDA() called without calling DMDATSSetIJacobian()");
116: PetscCall(DMGetLocalVector(dm, &Xloc));
117: PetscCall(DMGlobalToLocalBegin(dm, X, INSERT_VALUES, Xloc));
118: PetscCall(DMGlobalToLocalEnd(dm, X, INSERT_VALUES, Xloc));
119: PetscCall(DMGetLocalVector(dm, &Xdotloc));
120: PetscCall(DMGlobalToLocalBegin(dm, Xdot, INSERT_VALUES, Xdotloc));
121: PetscCall(DMGlobalToLocalEnd(dm, Xdot, INSERT_VALUES, Xdotloc));
122: PetscCall(DMDAGetLocalInfo(dm, &info));
123: PetscCall(DMDAVecGetArray(dm, Xloc, &x));
124: PetscCall(DMDAVecGetArray(dm, Xdotloc, &xdot));
125: CHKMEMQ;
126: PetscCall((*dmdats->ijacobianlocal)(&info, ptime, x, xdot, shift, A, B, dmdats->ijacobianlocalctx));
127: CHKMEMQ;
128: PetscCall(DMDAVecRestoreArray(dm, Xloc, &x));
129: PetscCall(DMDAVecRestoreArray(dm, Xdotloc, &xdot));
130: PetscCall(DMRestoreLocalVector(dm, &Xloc));
131: PetscCall(DMRestoreLocalVector(dm, &Xdotloc));
132: /* This will be redundant if the user called both, but it's too common to forget. */
133: if (A != B) {
134: PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
135: PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
136: }
137: PetscFunctionReturn(PETSC_SUCCESS);
138: }
140: static PetscErrorCode TSComputeRHSFunction_DMDA(TS ts, PetscReal ptime, Vec X, Vec F, void *ctx)
141: {
142: DM dm;
143: DMTS_DA *dmdats = (DMTS_DA *)ctx;
144: DMDALocalInfo info;
145: Vec Xloc;
146: void *x, *f;
148: PetscFunctionBegin;
152: PetscCheck(dmdats->rhsfunctionlocal, PetscObjectComm((PetscObject)ts), PETSC_ERR_PLIB, "Corrupt context");
153: PetscCall(TSGetDM(ts, &dm));
154: PetscCall(DMGetLocalVector(dm, &Xloc));
155: PetscCall(DMGlobalToLocalBegin(dm, X, INSERT_VALUES, Xloc));
156: PetscCall(DMGlobalToLocalEnd(dm, X, INSERT_VALUES, Xloc));
157: PetscCall(DMDAGetLocalInfo(dm, &info));
158: PetscCall(DMDAVecGetArray(dm, Xloc, &x));
159: switch (dmdats->rhsfunctionlocalimode) {
160: case INSERT_VALUES: {
161: PetscCall(DMDAVecGetArray(dm, F, &f));
162: CHKMEMQ;
163: PetscCall((*dmdats->rhsfunctionlocal)(&info, ptime, x, f, dmdats->rhsfunctionlocalctx));
164: CHKMEMQ;
165: PetscCall(DMDAVecRestoreArray(dm, F, &f));
166: } break;
167: case ADD_VALUES: {
168: Vec Floc;
169: PetscCall(DMGetLocalVector(dm, &Floc));
170: PetscCall(VecZeroEntries(Floc));
171: PetscCall(DMDAVecGetArray(dm, Floc, &f));
172: CHKMEMQ;
173: PetscCall((*dmdats->rhsfunctionlocal)(&info, ptime, x, f, dmdats->rhsfunctionlocalctx));
174: CHKMEMQ;
175: PetscCall(DMDAVecRestoreArray(dm, Floc, &f));
176: PetscCall(VecZeroEntries(F));
177: PetscCall(DMLocalToGlobalBegin(dm, Floc, ADD_VALUES, F));
178: PetscCall(DMLocalToGlobalEnd(dm, Floc, ADD_VALUES, F));
179: PetscCall(DMRestoreLocalVector(dm, &Floc));
180: } break;
181: default:
182: SETERRQ(PetscObjectComm((PetscObject)ts), PETSC_ERR_ARG_INCOMP, "Cannot use imode=%d", (int)dmdats->rhsfunctionlocalimode);
183: }
184: PetscCall(DMDAVecRestoreArray(dm, Xloc, &x));
185: PetscCall(DMRestoreLocalVector(dm, &Xloc));
186: PetscFunctionReturn(PETSC_SUCCESS);
187: }
189: static PetscErrorCode TSComputeRHSJacobian_DMDA(TS ts, PetscReal ptime, Vec X, Mat A, Mat B, void *ctx)
190: {
191: DM dm;
192: DMTS_DA *dmdats = (DMTS_DA *)ctx;
193: DMDALocalInfo info;
194: Vec Xloc;
195: void *x;
197: PetscFunctionBegin;
198: PetscCheck(dmdats->rhsfunctionlocal, PetscObjectComm((PetscObject)ts), PETSC_ERR_PLIB, "Corrupt context");
199: PetscCall(TSGetDM(ts, &dm));
201: PetscCheck(dmdats->rhsjacobianlocal, PetscObjectComm((PetscObject)ts), PETSC_ERR_PLIB, "TSComputeRHSJacobian_DMDA() called without calling DMDATSSetRHSJacobian()");
202: PetscCall(DMGetLocalVector(dm, &Xloc));
203: PetscCall(DMGlobalToLocalBegin(dm, X, INSERT_VALUES, Xloc));
204: PetscCall(DMGlobalToLocalEnd(dm, X, INSERT_VALUES, Xloc));
205: PetscCall(DMDAGetLocalInfo(dm, &info));
206: PetscCall(DMDAVecGetArray(dm, Xloc, &x));
207: CHKMEMQ;
208: PetscCall((*dmdats->rhsjacobianlocal)(&info, ptime, x, A, B, dmdats->rhsjacobianlocalctx));
209: CHKMEMQ;
210: PetscCall(DMDAVecRestoreArray(dm, Xloc, &x));
211: PetscCall(DMRestoreLocalVector(dm, &Xloc));
212: /* This will be redundant if the user called both, but it's too common to forget. */
213: if (A != B) {
214: PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
215: PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
216: }
217: PetscFunctionReturn(PETSC_SUCCESS);
218: }
220: /*@C
221: DMDATSSetRHSFunctionLocal - set a local residual evaluation function for use with `DMDA`
223: Logically Collective
225: Input Parameters:
226: + dm - `DM` to associate callback with
227: . imode - insert mode for the residual
228: . func - local residual evaluation, see `DMDATSRHSFunctionLocalFn` for the calling sequence
229: - ctx - optional context for local residual evaluation
231: Level: beginner
233: .seealso: [](ch_ts), `DMDA`, `DMDATSRHSFunctionLocalFn`, `TS`, `TSSetRHSFunction()`, `DMTSSetRHSFunction()`, `DMDATSSetRHSJacobianLocal()`, `DMDASNESSetFunctionLocal()`
234: @*/
235: PetscErrorCode DMDATSSetRHSFunctionLocal(DM dm, InsertMode imode, DMDATSRHSFunctionLocalFn *func, void *ctx)
236: {
237: DMTS sdm;
238: DMTS_DA *dmdats;
240: PetscFunctionBegin;
242: PetscCall(DMGetDMTSWrite(dm, &sdm));
243: PetscCall(DMDATSGetContext(dm, sdm, &dmdats));
244: dmdats->rhsfunctionlocalimode = imode;
245: dmdats->rhsfunctionlocal = func;
246: dmdats->rhsfunctionlocalctx = ctx;
247: PetscCall(DMTSSetRHSFunction(dm, TSComputeRHSFunction_DMDA, dmdats));
248: PetscFunctionReturn(PETSC_SUCCESS);
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, see `DMDATSRHSJacobianLocalFn` for the calling sequence
259: - ctx - optional context for local jacobian evaluation
261: Level: beginner
263: .seealso: [](ch_ts), `DMDA`, `DMDATSRHSJacobianLocalFn`, `DMTSSetRHSJacobian()`,
264: `DMDATSSetRHSFunctionLocal()`, `DMDASNESSetJacobianLocal()`
265: @*/
266: PetscErrorCode DMDATSSetRHSJacobianLocal(DM dm, DMDATSRHSJacobianLocalFn *func, void *ctx)
267: {
268: DMTS sdm;
269: DMTS_DA *dmdats;
271: PetscFunctionBegin;
273: PetscCall(DMGetDMTSWrite(dm, &sdm));
274: PetscCall(DMDATSGetContext(dm, sdm, &dmdats));
275: dmdats->rhsjacobianlocal = func;
276: dmdats->rhsjacobianlocalctx = ctx;
277: PetscCall(DMTSSetRHSJacobian(dm, TSComputeRHSJacobian_DMDA, dmdats));
278: PetscFunctionReturn(PETSC_SUCCESS);
279: }
281: /*@C
282: DMDATSSetIFunctionLocal - set a local residual evaluation function for use with `DMDA`
284: Logically Collective
286: Input Parameters:
287: + dm - `DM` to associate callback with
288: . imode - the insert mode of the function
289: . func - local residual evaluation, see `DMDATSIFunctionLocalFn` for the calling sequence
290: - ctx - optional context for local residual evaluation
292: Level: beginner
294: .seealso: [](ch_ts), `DMDA`, `DMDATSIFunctionLocalFn`, `DMTSSetIFunction()`,
295: `DMDATSSetIJacobianLocal()`, `DMDASNESSetFunctionLocal()`
296: @*/
297: PetscErrorCode DMDATSSetIFunctionLocal(DM dm, InsertMode imode, DMDATSIFunctionLocalFn *func, void *ctx)
298: {
299: DMTS sdm;
300: DMTS_DA *dmdats;
302: PetscFunctionBegin;
304: PetscCall(DMGetDMTSWrite(dm, &sdm));
305: PetscCall(DMDATSGetContext(dm, sdm, &dmdats));
306: dmdats->ifunctionlocalimode = imode;
307: dmdats->ifunctionlocal = func;
308: dmdats->ifunctionlocalctx = ctx;
309: PetscCall(DMTSSetIFunction(dm, TSComputeIFunction_DMDA, dmdats));
310: PetscFunctionReturn(PETSC_SUCCESS);
311: }
313: /*@C
314: DMDATSSetIJacobianLocal - set a local residual evaluation function for use with `DMDA`
316: Logically Collective
318: Input Parameters:
319: + dm - `DM` to associate callback with
320: . func - local residual evaluation, see `DMDATSIJacobianLocalFn` for the calling sequence
321: - ctx - optional context for local residual evaluation
323: Level: beginner
325: .seealso: [](ch_ts), `DMDA`, `DMDATSIJacobianLocalFn`, `DMTSSetIJacobian()`,
326: `DMDATSSetIFunctionLocal()`, `DMDASNESSetJacobianLocal()`
327: @*/
328: PetscErrorCode DMDATSSetIJacobianLocal(DM dm, DMDATSIJacobianLocalFn *func, void *ctx)
329: {
330: DMTS sdm;
331: DMTS_DA *dmdats;
333: PetscFunctionBegin;
335: PetscCall(DMGetDMTSWrite(dm, &sdm));
336: PetscCall(DMDATSGetContext(dm, sdm, &dmdats));
337: dmdats->ijacobianlocal = func;
338: dmdats->ijacobianlocalctx = ctx;
339: PetscCall(DMTSSetIJacobian(dm, TSComputeIJacobian_DMDA, dmdats));
340: PetscFunctionReturn(PETSC_SUCCESS);
341: }
343: PetscErrorCode TSMonitorDMDARayDestroy(void **mctx)
344: {
345: TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx *)*mctx;
347: PetscFunctionBegin;
348: if (rayctx->lgctx) PetscCall(TSMonitorLGCtxDestroy(&rayctx->lgctx));
349: PetscCall(VecDestroy(&rayctx->ray));
350: PetscCall(VecScatterDestroy(&rayctx->scatter));
351: PetscCall(PetscViewerDestroy(&rayctx->viewer));
352: PetscCall(PetscFree(rayctx));
353: PetscFunctionReturn(PETSC_SUCCESS);
354: }
356: PetscErrorCode TSMonitorDMDARay(TS ts, PetscInt steps, PetscReal time, Vec u, void *mctx)
357: {
358: TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx *)mctx;
359: Vec solution;
361: PetscFunctionBegin;
362: PetscCall(TSGetSolution(ts, &solution));
363: PetscCall(VecScatterBegin(rayctx->scatter, solution, rayctx->ray, INSERT_VALUES, SCATTER_FORWARD));
364: PetscCall(VecScatterEnd(rayctx->scatter, solution, rayctx->ray, INSERT_VALUES, SCATTER_FORWARD));
365: if (rayctx->viewer) PetscCall(VecView(rayctx->ray, rayctx->viewer));
366: PetscFunctionReturn(PETSC_SUCCESS);
367: }
369: PetscErrorCode TSMonitorLGDMDARay(TS ts, PetscInt step, PetscReal ptime, Vec u, void *ctx)
370: {
371: TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx *)ctx;
372: TSMonitorLGCtx lgctx = rayctx->lgctx;
373: Vec v = rayctx->ray;
374: const PetscScalar *a;
375: PetscInt dim;
377: PetscFunctionBegin;
378: PetscCall(VecScatterBegin(rayctx->scatter, u, v, INSERT_VALUES, SCATTER_FORWARD));
379: PetscCall(VecScatterEnd(rayctx->scatter, u, v, INSERT_VALUES, SCATTER_FORWARD));
380: if (!step) {
381: PetscDrawAxis axis;
383: PetscCall(PetscDrawLGGetAxis(lgctx->lg, &axis));
384: PetscCall(PetscDrawAxisSetLabels(axis, "Solution Ray as function of time", "Time", "Solution"));
385: PetscCall(VecGetLocalSize(rayctx->ray, &dim));
386: PetscCall(PetscDrawLGSetDimension(lgctx->lg, dim));
387: PetscCall(PetscDrawLGReset(lgctx->lg));
388: }
389: PetscCall(VecGetArrayRead(v, &a));
390: #if defined(PETSC_USE_COMPLEX)
391: {
392: PetscReal *areal;
393: PetscInt i, n;
394: PetscCall(VecGetLocalSize(v, &n));
395: PetscCall(PetscMalloc1(n, &areal));
396: for (i = 0; i < n; ++i) areal[i] = PetscRealPart(a[i]);
397: PetscCall(PetscDrawLGAddCommonPoint(lgctx->lg, ptime, areal));
398: PetscCall(PetscFree(areal));
399: }
400: #else
401: PetscCall(PetscDrawLGAddCommonPoint(lgctx->lg, ptime, a));
402: #endif
403: PetscCall(VecRestoreArrayRead(v, &a));
404: if (((lgctx->howoften > 0) && (!(step % lgctx->howoften))) || ((lgctx->howoften == -1) && ts->reason)) {
405: PetscCall(PetscDrawLGDraw(lgctx->lg));
406: PetscCall(PetscDrawLGSave(lgctx->lg));
407: }
408: PetscFunctionReturn(PETSC_SUCCESS);
409: }