xref: /petsc/src/ts/utils/dmdats.c (revision 6a5217c03994f2d95bb2e6dbd8bed42381aeb015)
1 #include <petscdmda.h>          /*I "petscdmda.h" I*/
2 #include <petsc/private/dmimpl.h>
3 #include <petsc/private/tsimpl.h>   /*I "petscts.h" I*/
4 #include <petscdraw.h>
5 
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;
19 
20 static PetscErrorCode DMTSDestroy_DMDA(DMTS sdm)
21 {
22   PetscFunctionBegin;
23   PetscCall(PetscFree(sdm->data));
24   PetscFunctionReturn(0);
25 }
26 
27 static PetscErrorCode DMTSDuplicate_DMDA(DMTS oldsdm,DMTS sdm)
28 {
29   PetscFunctionBegin;
30   PetscCall(PetscNewLog(sdm,(DMTS_DA**)&sdm->data));
31   if (oldsdm->data) PetscCall(PetscMemcpy(sdm->data,oldsdm->data,sizeof(DMTS_DA)));
32   PetscFunctionReturn(0);
33 }
34 
35 static PetscErrorCode DMDATSGetContext(DM dm,DMTS sdm,DMTS_DA **dmdats)
36 {
37   PetscFunctionBegin;
38   *dmdats = NULL;
39   if (!sdm->data) {
40     PetscCall(PetscNewLog(dm,(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(0);
46 }
47 
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;
55 
56   PetscFunctionBegin;
57   PetscValidHeaderSpecific(ts,TS_CLASSID,1);
58   PetscValidHeaderSpecific(X,VEC_CLASSID,3);
59   PetscValidHeaderSpecific(F,VEC_CLASSID,5);
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: SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_ARG_INCOMP,"Cannot use imode=%d",(int)dmdats->ifunctionlocalimode);
94   }
95   PetscCall(DMDAVecRestoreArray(dm,Xloc,&x));
96   PetscCall(DMRestoreLocalVector(dm,&Xloc));
97   PetscCall(DMDAVecRestoreArray(dm,Xdotloc,&xdot));
98   PetscCall(DMRestoreLocalVector(dm,&Xdotloc));
99   PetscFunctionReturn(0);
100 }
101 
102 static PetscErrorCode TSComputeIJacobian_DMDA(TS ts,PetscReal ptime,Vec X,Vec Xdot,PetscReal shift,Mat A,Mat B,void *ctx)
103 {
104   DM             dm;
105   DMTS_DA        *dmdats = (DMTS_DA*)ctx;
106   DMDALocalInfo  info;
107   Vec            Xloc;
108   void           *x,*xdot;
109 
110   PetscFunctionBegin;
111   PetscCheck(dmdats->ifunctionlocal,PetscObjectComm((PetscObject)ts),PETSC_ERR_PLIB,"Corrupt context");
112   PetscCall(TSGetDM(ts,&dm));
113 
114   if (dmdats->ijacobianlocal) {
115     PetscCall(DMGetLocalVector(dm,&Xloc));
116     PetscCall(DMGlobalToLocalBegin(dm,X,INSERT_VALUES,Xloc));
117     PetscCall(DMGlobalToLocalEnd(dm,X,INSERT_VALUES,Xloc));
118     PetscCall(DMDAGetLocalInfo(dm,&info));
119     PetscCall(DMDAVecGetArray(dm,Xloc,&x));
120     PetscCall(DMDAVecGetArray(dm,Xdot,&xdot));
121     CHKMEMQ;
122     PetscCall((*dmdats->ijacobianlocal)(&info,ptime,x,xdot,shift,A,B,dmdats->ijacobianlocalctx));
123     CHKMEMQ;
124     PetscCall(DMDAVecRestoreArray(dm,Xloc,&x));
125     PetscCall(DMDAVecRestoreArray(dm,Xdot,&xdot));
126     PetscCall(DMRestoreLocalVector(dm,&Xloc));
127   } else SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_PLIB,"TSComputeIJacobian_DMDA() called without calling DMDATSSetIJacobian()");
128   /* This will be redundant if the user called both, but it's too common to forget. */
129   if (A != B) {
130     PetscCall(MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY));
131     PetscCall(MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY));
132   }
133   PetscFunctionReturn(0);
134 }
135 
136 static PetscErrorCode TSComputeRHSFunction_DMDA(TS ts,PetscReal ptime,Vec X,Vec F,void *ctx)
137 {
138   DM             dm;
139   DMTS_DA        *dmdats = (DMTS_DA*)ctx;
140   DMDALocalInfo  info;
141   Vec            Xloc;
142   void           *x,*f;
143 
144   PetscFunctionBegin;
145   PetscValidHeaderSpecific(ts,TS_CLASSID,1);
146   PetscValidHeaderSpecific(X,VEC_CLASSID,3);
147   PetscValidHeaderSpecific(F,VEC_CLASSID,4);
148   PetscCheck(dmdats->rhsfunctionlocal,PetscObjectComm((PetscObject)ts),PETSC_ERR_PLIB,"Corrupt context");
149   PetscCall(TSGetDM(ts,&dm));
150   PetscCall(DMGetLocalVector(dm,&Xloc));
151   PetscCall(DMGlobalToLocalBegin(dm,X,INSERT_VALUES,Xloc));
152   PetscCall(DMGlobalToLocalEnd(dm,X,INSERT_VALUES,Xloc));
153   PetscCall(DMDAGetLocalInfo(dm,&info));
154   PetscCall(DMDAVecGetArray(dm,Xloc,&x));
155   switch (dmdats->rhsfunctionlocalimode) {
156   case INSERT_VALUES: {
157     PetscCall(DMDAVecGetArray(dm,F,&f));
158     CHKMEMQ;
159     PetscCall((*dmdats->rhsfunctionlocal)(&info,ptime,x,f,dmdats->rhsfunctionlocalctx));
160     CHKMEMQ;
161     PetscCall(DMDAVecRestoreArray(dm,F,&f));
162   } break;
163   case ADD_VALUES: {
164     Vec Floc;
165     PetscCall(DMGetLocalVector(dm,&Floc));
166     PetscCall(VecZeroEntries(Floc));
167     PetscCall(DMDAVecGetArray(dm,Floc,&f));
168     CHKMEMQ;
169     PetscCall((*dmdats->rhsfunctionlocal)(&info,ptime,x,f,dmdats->rhsfunctionlocalctx));
170     CHKMEMQ;
171     PetscCall(DMDAVecRestoreArray(dm,Floc,&f));
172     PetscCall(VecZeroEntries(F));
173     PetscCall(DMLocalToGlobalBegin(dm,Floc,ADD_VALUES,F));
174     PetscCall(DMLocalToGlobalEnd(dm,Floc,ADD_VALUES,F));
175     PetscCall(DMRestoreLocalVector(dm,&Floc));
176   } break;
177   default: SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_ARG_INCOMP,"Cannot use imode=%d",(int)dmdats->rhsfunctionlocalimode);
178   }
179   PetscCall(DMDAVecRestoreArray(dm,Xloc,&x));
180   PetscCall(DMRestoreLocalVector(dm,&Xloc));
181   PetscFunctionReturn(0);
182 }
183 
184 static PetscErrorCode TSComputeRHSJacobian_DMDA(TS ts,PetscReal ptime,Vec X,Mat A,Mat B,void *ctx)
185 {
186   DM             dm;
187   DMTS_DA        *dmdats = (DMTS_DA*)ctx;
188   DMDALocalInfo  info;
189   Vec            Xloc;
190   void           *x;
191 
192   PetscFunctionBegin;
193   PetscCheck(dmdats->rhsfunctionlocal,PetscObjectComm((PetscObject)ts),PETSC_ERR_PLIB,"Corrupt context");
194   PetscCall(TSGetDM(ts,&dm));
195 
196   if (dmdats->rhsjacobianlocal) {
197     PetscCall(DMGetLocalVector(dm,&Xloc));
198     PetscCall(DMGlobalToLocalBegin(dm,X,INSERT_VALUES,Xloc));
199     PetscCall(DMGlobalToLocalEnd(dm,X,INSERT_VALUES,Xloc));
200     PetscCall(DMDAGetLocalInfo(dm,&info));
201     PetscCall(DMDAVecGetArray(dm,Xloc,&x));
202     CHKMEMQ;
203     PetscCall((*dmdats->rhsjacobianlocal)(&info,ptime,x,A,B,dmdats->rhsjacobianlocalctx));
204     CHKMEMQ;
205     PetscCall(DMDAVecRestoreArray(dm,Xloc,&x));
206     PetscCall(DMRestoreLocalVector(dm,&Xloc));
207   } else SETERRQ(PetscObjectComm((PetscObject)ts),PETSC_ERR_PLIB,"TSComputeRHSJacobian_DMDA() called without calling DMDATSSetRHSJacobian()");
208   /* This will be redundant if the user called both, but it's too common to forget. */
209   if (A != B) {
210     PetscCall(MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY));
211     PetscCall(MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY));
212   }
213   PetscFunctionReturn(0);
214 }
215 
216 /*@C
217    DMDATSSetRHSFunctionLocal - set a local residual evaluation function
218 
219    Logically Collective
220 
221    Input Parameters:
222 +  dm - DM to associate callback with
223 .  imode - insert mode for the residual
224 .  func - local residual evaluation
225 -  ctx - optional context for local residual evaluation
226 
227    Calling sequence for func:
228 
229 $ func(DMDALocalInfo info,PetscReal t,void *x,void *f,void *ctx)
230 
231 +  info - DMDALocalInfo defining 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
236 
237    Level: beginner
238 
239 .seealso: DMTSSetRHSFunction(), DMDATSSetRHSJacobianLocal(), DMDASNESSetFunctionLocal()
240 @*/
241 PetscErrorCode DMDATSSetRHSFunctionLocal(DM dm,InsertMode imode,DMDATSRHSFunctionLocal func,void *ctx)
242 {
243   DMTS           sdm;
244   DMTS_DA        *dmdats;
245 
246   PetscFunctionBegin;
247   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
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(0);
255 }
256 
257 /*@C
258    DMDATSSetRHSJacobianLocal - set a local residual evaluation function
259 
260    Logically Collective
261 
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
266 
267    Calling sequence for func:
268 
269 $ func(DMDALocalInfo* info,PetscReal t,void* x,Mat J,Mat B,void *ctx);
270 
271 +  info - DMDALocalInfo defining the subdomain to evaluate the residual on
272 .  t    - time at which to evaluate residual
273 .  x    - array of local state information
274 .  J    - Jacobian matrix
275 .  B    - preconditioner matrix; often same as J
276 -  ctx  - optional context passed above
277 
278    Level: beginner
279 
280 .seealso: DMTSSetRHSJacobian(), DMDATSSetRHSFunctionLocal(), DMDASNESSetJacobianLocal()
281 @*/
282 PetscErrorCode DMDATSSetRHSJacobianLocal(DM dm,DMDATSRHSJacobianLocal func,void *ctx)
283 {
284   DMTS           sdm;
285   DMTS_DA        *dmdats;
286 
287   PetscFunctionBegin;
288   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
289   PetscCall(DMGetDMTSWrite(dm,&sdm));
290   PetscCall(DMDATSGetContext(dm,sdm,&dmdats));
291   dmdats->rhsjacobianlocal    = func;
292   dmdats->rhsjacobianlocalctx = ctx;
293   PetscCall(DMTSSetRHSJacobian(dm,TSComputeRHSJacobian_DMDA,dmdats));
294   PetscFunctionReturn(0);
295 }
296 
297 /*@C
298    DMDATSSetIFunctionLocal - set a local residual evaluation function
299 
300    Logically Collective
301 
302    Input Parameters:
303 +  dm   - DM to associate callback with
304 .  func - local residual evaluation
305 -  ctx  - optional context for local residual evaluation
306 
307    Calling sequence for func:
308 +  info - DMDALocalInfo defining the subdomain to evaluate the residual on
309 .  t    - time at which to evaluate residual
310 .  x    - array of local state information
311 .  xdot - array of local time derivative information
312 .  f    - output array of local function evaluation information
313 -  ctx - optional context passed above
314 
315    Level: beginner
316 
317 .seealso: DMTSSetIFunction(), DMDATSSetIJacobianLocal(), DMDASNESSetFunctionLocal()
318 @*/
319 PetscErrorCode DMDATSSetIFunctionLocal(DM dm,InsertMode imode,DMDATSIFunctionLocal func,void *ctx)
320 {
321   DMTS           sdm;
322   DMTS_DA        *dmdats;
323 
324   PetscFunctionBegin;
325   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
326   PetscCall(DMGetDMTSWrite(dm,&sdm));
327   PetscCall(DMDATSGetContext(dm,sdm,&dmdats));
328   dmdats->ifunctionlocalimode = imode;
329   dmdats->ifunctionlocal      = func;
330   dmdats->ifunctionlocalctx   = ctx;
331   PetscCall(DMTSSetIFunction(dm,TSComputeIFunction_DMDA,dmdats));
332   PetscFunctionReturn(0);
333 }
334 
335 /*@C
336    DMDATSSetIJacobianLocal - set a local residual evaluation function
337 
338    Logically Collective
339 
340    Input Parameters:
341 +  dm   - DM to associate callback with
342 .  func - local residual evaluation
343 -  ctx   - optional context for local residual evaluation
344 
345    Calling sequence for func:
346 
347 $ func(DMDALocalInfo* info,PetscReal t,void* x,void *xdot,PetscScalar shift,Mat J,Mat B,void *ctx);
348 
349 +  info - DMDALocalInfo defining the subdomain to evaluate the residual on
350 .  t    - time at which to evaluate the jacobian
351 .  x    - array of local state information
352 .  xdot - time derivative at this state
353 .  shift - see TSSetIJacobian() for the meaning of this parameter
354 .  J    - Jacobian matrix
355 .  B    - preconditioner matrix; often same as J
356 -  ctx  - optional context passed above
357 
358    Level: beginner
359 
360 .seealso: DMTSSetJacobian(), DMDATSSetIFunctionLocal(), DMDASNESSetJacobianLocal()
361 @*/
362 PetscErrorCode DMDATSSetIJacobianLocal(DM dm,DMDATSIJacobianLocal func,void *ctx)
363 {
364   DMTS           sdm;
365   DMTS_DA        *dmdats;
366 
367   PetscFunctionBegin;
368   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
369   PetscCall(DMGetDMTSWrite(dm,&sdm));
370   PetscCall(DMDATSGetContext(dm,sdm,&dmdats));
371   dmdats->ijacobianlocal    = func;
372   dmdats->ijacobianlocalctx = ctx;
373   PetscCall(DMTSSetIJacobian(dm,TSComputeIJacobian_DMDA,dmdats));
374   PetscFunctionReturn(0);
375 }
376 
377 PetscErrorCode TSMonitorDMDARayDestroy(void **mctx)
378 {
379   TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx *) *mctx;
380 
381   PetscFunctionBegin;
382   if (rayctx->lgctx) PetscCall(TSMonitorLGCtxDestroy(&rayctx->lgctx));
383   PetscCall(VecDestroy(&rayctx->ray));
384   PetscCall(VecScatterDestroy(&rayctx->scatter));
385   PetscCall(PetscViewerDestroy(&rayctx->viewer));
386   PetscCall(PetscFree(rayctx));
387   PetscFunctionReturn(0);
388 }
389 
390 PetscErrorCode TSMonitorDMDARay(TS ts,PetscInt steps,PetscReal time,Vec u,void *mctx)
391 {
392   TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx*)mctx;
393   Vec                 solution;
394 
395   PetscFunctionBegin;
396   PetscCall(TSGetSolution(ts,&solution));
397   PetscCall(VecScatterBegin(rayctx->scatter,solution,rayctx->ray,INSERT_VALUES,SCATTER_FORWARD));
398   PetscCall(VecScatterEnd(rayctx->scatter,solution,rayctx->ray,INSERT_VALUES,SCATTER_FORWARD));
399   if (rayctx->viewer) {
400     PetscCall(VecView(rayctx->ray,rayctx->viewer));
401   }
402   PetscFunctionReturn(0);
403 }
404 
405 PetscErrorCode  TSMonitorLGDMDARay(TS ts, PetscInt step, PetscReal ptime, Vec u, void *ctx)
406 {
407   TSMonitorDMDARayCtx *rayctx = (TSMonitorDMDARayCtx *) ctx;
408   TSMonitorLGCtx       lgctx  = (TSMonitorLGCtx) rayctx->lgctx;
409   Vec                  v      = rayctx->ray;
410   const PetscScalar   *a;
411   PetscInt             dim;
412 
413   PetscFunctionBegin;
414   PetscCall(VecScatterBegin(rayctx->scatter, u, v, INSERT_VALUES, SCATTER_FORWARD));
415   PetscCall(VecScatterEnd(rayctx->scatter, u, v, INSERT_VALUES, SCATTER_FORWARD));
416   if (!step) {
417     PetscDrawAxis axis;
418 
419     PetscCall(PetscDrawLGGetAxis(lgctx->lg, &axis));
420     PetscCall(PetscDrawAxisSetLabels(axis, "Solution Ray as function of time", "Time", "Solution"));
421     PetscCall(VecGetLocalSize(rayctx->ray, &dim));
422     PetscCall(PetscDrawLGSetDimension(lgctx->lg, dim));
423     PetscCall(PetscDrawLGReset(lgctx->lg));
424   }
425   PetscCall(VecGetArrayRead(v, &a));
426 #if defined(PETSC_USE_COMPLEX)
427   {
428     PetscReal *areal;
429     PetscInt   i,n;
430     PetscCall(VecGetLocalSize(v, &n));
431     PetscCall(PetscMalloc1(n, &areal));
432     for (i = 0; i < n; ++i) areal[i] = PetscRealPart(a[i]);
433     PetscCall(PetscDrawLGAddCommonPoint(lgctx->lg, ptime, areal));
434     PetscCall(PetscFree(areal));
435   }
436 #else
437   PetscCall(PetscDrawLGAddCommonPoint(lgctx->lg, ptime, a));
438 #endif
439   PetscCall(VecRestoreArrayRead(v, &a));
440   if (((lgctx->howoften > 0) && (!(step % lgctx->howoften))) || ((lgctx->howoften == -1) && ts->reason)) {
441     PetscCall(PetscDrawLGDraw(lgctx->lg));
442     PetscCall(PetscDrawLGSave(lgctx->lg));
443   }
444   PetscFunctionReturn(0);
445 }
446