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