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: Level: beginner
231: .seealso: [](ch_ts), `DMDA`, `DMDATSRHSFunctionLocal`, `TS`, `TSSetRHSFunction()`, `DMTSSetRHSFunction()`, `DMDATSSetRHSJacobianLocal()`, `DMDASNESSetFunctionLocal()`
232: @*/
233: PetscErrorCode DMDATSSetRHSFunctionLocal(DM dm, InsertMode imode, DMDATSRHSFunctionLocal 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
257: - ctx - optional context for local jacobian evaluation
259: Level: beginner
261: .seealso: [](ch_ts), `DMDA`, `DMDATSRHSJacobianLocal`, `DMTSSetRHSJacobian()`,
262: `DMDATSSetRHSFunctionLocal()`, `DMDASNESSetJacobianLocal()`
263: @*/
264: PetscErrorCode DMDATSSetRHSJacobianLocal(DM dm, DMDATSRHSJacobianLocal 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
288: - ctx - optional context for local residual evaluation
290: Level: beginner
292: .seealso: [](ch_ts), `DMDA`, `DMDATSIFunctionLocal`, `DMTSSetIFunction()`,
293: `DMDATSSetIJacobianLocal()`, `DMDASNESSetFunctionLocal()`
294: @*/
295: PetscErrorCode DMDATSSetIFunctionLocal(DM dm, InsertMode imode, DMDATSIFunctionLocal 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
319: - ctx - optional context for local residual evaluation
321: Level: beginner
323: .seealso: [](ch_ts), `DMDA`, `DMDATSIJacobianLocal`, `DMTSSetJacobian()`,
324: `DMDATSSetIFunctionLocal()`, `DMDASNESSetJacobianLocal()`
325: @*/
326: PetscErrorCode DMDATSSetIJacobianLocal(DM dm, DMDATSIJacobianLocal 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: }