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: }