xref: /petsc/src/ts/utils/dmdats.c (revision 3923b477fd0dced8a2d147b4fb4519fe3af97d3f)
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 
5 /* This structure holds the user-provided DMDA callbacks */
6 typedef struct {
7   PetscErrorCode (*ifunctionlocal)(DMDALocalInfo*,PetscReal,void*,void*,void*,void*);
8   PetscErrorCode (*rhsfunctionlocal)(DMDALocalInfo*,PetscReal,void*,void*,void*);
9   PetscErrorCode (*ijacobianlocal)(DMDALocalInfo*,PetscReal,void*,void*,PetscReal,Mat,Mat,MatStructure*,void*);
10   PetscErrorCode (*rhsjacobianlocal)(DMDALocalInfo*,PetscReal,void*,Mat,Mat,MatStructure*,void*);
11   void *ifunctionlocalctx;
12   void *ijacobianlocalctx;
13   void *rhsfunctionlocalctx;
14   void *rhsjacobianlocalctx;
15   InsertMode ifunctionlocalimode;
16   InsertMode rhsfunctionlocalimode;
17 } DM_DA_TS;
18 
19 #undef __FUNCT__
20 #define __FUNCT__ "TSDMDestroy_DMDA"
21 static PetscErrorCode TSDMDestroy_DMDA(TSDM sdm)
22 {
23   PetscErrorCode ierr;
24 
25   PetscFunctionBegin;
26   ierr = PetscFree(sdm->data);CHKERRQ(ierr);
27   PetscFunctionReturn(0);
28 }
29 
30 #undef __FUNCT__
31 #define __FUNCT__ "DMDATSGetContext"
32 static PetscErrorCode DMDATSGetContext(DM dm,TSDM sdm,DM_DA_TS **dmdats)
33 {
34   PetscErrorCode ierr;
35 
36   PetscFunctionBegin;
37   *dmdats = PETSC_NULL;
38   if (!sdm->data) {
39     ierr = PetscNewLog(dm,DM_DA_TS,&sdm->data);CHKERRQ(ierr);
40     sdm->destroy = TSDMDestroy_DMDA;
41   }
42   *dmdats = (DM_DA_TS*)sdm->data;
43   PetscFunctionReturn(0);
44 }
45 
46 #undef __FUNCT__
47 #define __FUNCT__ "TSComputeIFunction_DMDA"
48 /*
49   This function should eventually replace:
50     DMDAComputeFunction() and DMDAComputeFunction1()
51  */
52 static PetscErrorCode TSComputeIFunction_DMDA(TS ts,PetscReal ptime,Vec X,Vec Xdot,Vec F,void *ctx)
53 {
54   PetscErrorCode ierr;
55   DM             dm;
56   DM_DA_TS     *dmdats = (DM_DA_TS*)ctx;
57   DMDALocalInfo  info;
58   Vec            Xloc;
59   void           *x,*f,*xdot;
60 
61   PetscFunctionBegin;
62   PetscValidHeaderSpecific(ts,TS_CLASSID,1);
63   PetscValidHeaderSpecific(X,VEC_CLASSID,2);
64   PetscValidHeaderSpecific(F,VEC_CLASSID,3);
65   if (!dmdats->ifunctionlocal) SETERRQ(((PetscObject)ts)->comm,PETSC_ERR_PLIB,"Corrupt context");
66   ierr = TSGetDM(ts,&dm);CHKERRQ(ierr);
67   ierr = DMGetLocalVector(dm,&Xloc);CHKERRQ(ierr);
68   ierr = DMGlobalToLocalBegin(dm,X,INSERT_VALUES,Xloc);CHKERRQ(ierr);
69   ierr = DMGlobalToLocalEnd(dm,X,INSERT_VALUES,Xloc);CHKERRQ(ierr);
70   ierr = DMDAGetLocalInfo(dm,&info);CHKERRQ(ierr);
71   ierr = DMDAVecGetArray(dm,Xloc,&x);CHKERRQ(ierr);
72   ierr = DMDAVecGetArray(dm,Xdot,&xdot);CHKERRQ(ierr);
73   switch (dmdats->ifunctionlocalimode) {
74   case INSERT_VALUES: {
75     ierr = DMDAVecGetArray(dm,F,&f);CHKERRQ(ierr);
76     CHKMEMQ;
77     ierr = (*dmdats->ifunctionlocal)(&info,ptime,x,xdot,f,dmdats->ifunctionlocalctx);CHKERRQ(ierr);
78     CHKMEMQ;
79     ierr = DMDAVecRestoreArray(dm,F,&f);CHKERRQ(ierr);
80   } break;
81   case ADD_VALUES: {
82     Vec Floc;
83     ierr = DMGetLocalVector(dm,&Floc);CHKERRQ(ierr);
84     ierr = VecZeroEntries(Floc);CHKERRQ(ierr);
85     ierr = DMDAVecGetArray(dm,Floc,&f);CHKERRQ(ierr);
86     CHKMEMQ;
87     ierr = (*dmdats->ifunctionlocal)(&info,ptime,x,xdot,f,dmdats->ifunctionlocalctx);CHKERRQ(ierr);
88     CHKMEMQ;
89     ierr = DMDAVecRestoreArray(dm,Floc,&f);CHKERRQ(ierr);
90     ierr = VecZeroEntries(F);CHKERRQ(ierr);
91     ierr = DMLocalToGlobalBegin(dm,Floc,ADD_VALUES,F);CHKERRQ(ierr);
92     ierr = DMLocalToGlobalEnd(dm,Floc,ADD_VALUES,F);CHKERRQ(ierr);
93     ierr = DMRestoreLocalVector(dm,&Floc);CHKERRQ(ierr);
94   } break;
95   default: SETERRQ1(((PetscObject)ts)->comm,PETSC_ERR_ARG_INCOMP,"Cannot use imode=%d",(int)dmdats->ifunctionlocalimode);
96   }
97   ierr = DMDAVecRestoreArray(dm,Xloc,&x);CHKERRQ(ierr);
98   ierr = DMRestoreLocalVector(dm,&Xloc);CHKERRQ(ierr);
99   PetscFunctionReturn(0);
100 }
101 
102 #undef __FUNCT__
103 #define __FUNCT__ "TSComputeIJacobian_DMDA"
104 /*
105   This function should eventually replace:
106     DMComputeJacobian() and DMDAComputeJacobian1()
107  */
108 static PetscErrorCode TSComputeIJacobian_DMDA(TS ts,PetscReal ptime,Vec X,Vec Xdot,PetscReal shift,Mat *A,Mat *B,MatStructure *mstr,void *ctx)
109 {
110   PetscErrorCode ierr;
111   DM             dm;
112   DM_DA_TS     *dmdats = (DM_DA_TS*)ctx;
113   DMDALocalInfo  info;
114   Vec            Xloc;
115   void           *x,*xdot;
116 
117   PetscFunctionBegin;
118   if (!dmdats->ifunctionlocal) SETERRQ(((PetscObject)ts)->comm,PETSC_ERR_PLIB,"Corrupt context");
119   ierr = TSGetDM(ts,&dm);CHKERRQ(ierr);
120 
121   if (dmdats->ijacobianlocal) {
122     ierr = DMGetLocalVector(dm,&Xloc);CHKERRQ(ierr);
123     ierr = DMGlobalToLocalBegin(dm,X,INSERT_VALUES,Xloc);CHKERRQ(ierr);
124     ierr = DMGlobalToLocalEnd(dm,X,INSERT_VALUES,Xloc);CHKERRQ(ierr);
125     ierr = DMDAGetLocalInfo(dm,&info);CHKERRQ(ierr);
126     ierr = DMDAVecGetArray(dm,Xloc,&x);CHKERRQ(ierr);
127     ierr = DMDAVecGetArray(dm,Xdot,&xdot);CHKERRQ(ierr);
128     CHKMEMQ;
129     ierr = (*dmdats->ijacobianlocal)(&info,ptime,x,xdot,shift,*A,*B,mstr,dmdats->ijacobianlocalctx);CHKERRQ(ierr);
130     CHKMEMQ;
131     ierr = DMDAVecRestoreArray(dm,Xloc,&x);CHKERRQ(ierr);
132     ierr = DMDAVecRestoreArray(dm,Xdot,&xdot);CHKERRQ(ierr);
133     ierr = DMRestoreLocalVector(dm,&Xloc);CHKERRQ(ierr);
134   } else {
135     SETERRQ(((PetscObject)ts)->comm,PETSC_ERR_PLIB,"TSComputeIJacobian_DMDA() called without calling DMDATSSetIJacobian()");
136   }
137   /* This will be redundant if the user called both, but it's too common to forget. */
138   if (*A != *B) {
139     ierr = MatAssemblyBegin(*A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
140     ierr = MatAssemblyEnd(*A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
141   }
142   PetscFunctionReturn(0);
143 }
144 
145 #undef __FUNCT__
146 #define __FUNCT__ "TSComputeRHSFunction_DMDA"
147 /*
148   This function should eventually replace:
149     DMDAComputeFunction() and DMDAComputeFunction1()
150  */
151 static PetscErrorCode TSComputeRHSFunction_DMDA(TS ts,PetscReal ptime,Vec X,Vec F,void *ctx)
152 {
153   PetscErrorCode ierr;
154   DM             dm;
155   DM_DA_TS     *dmdats = (DM_DA_TS*)ctx;
156   DMDALocalInfo  info;
157   Vec            Xloc;
158   void           *x,*f;
159 
160   PetscFunctionBegin;
161   PetscValidHeaderSpecific(ts,TS_CLASSID,1);
162   PetscValidHeaderSpecific(X,VEC_CLASSID,2);
163   PetscValidHeaderSpecific(F,VEC_CLASSID,3);
164   if (!dmdats->rhsfunctionlocal) SETERRQ(((PetscObject)ts)->comm,PETSC_ERR_PLIB,"Corrupt context");
165   ierr = TSGetDM(ts,&dm);CHKERRQ(ierr);
166   ierr = DMGetLocalVector(dm,&Xloc);CHKERRQ(ierr);
167   ierr = DMGlobalToLocalBegin(dm,X,INSERT_VALUES,Xloc);CHKERRQ(ierr);
168   ierr = DMGlobalToLocalEnd(dm,X,INSERT_VALUES,Xloc);CHKERRQ(ierr);
169   ierr = DMDAGetLocalInfo(dm,&info);CHKERRQ(ierr);
170   ierr = DMDAVecGetArray(dm,Xloc,&x);CHKERRQ(ierr);
171   switch (dmdats->rhsfunctionlocalimode) {
172   case INSERT_VALUES: {
173     ierr = DMDAVecGetArray(dm,F,&f);CHKERRQ(ierr);
174     CHKMEMQ;
175     ierr = (*dmdats->rhsfunctionlocal)(&info,ptime,x,f,dmdats->rhsfunctionlocalctx);CHKERRQ(ierr);
176     CHKMEMQ;
177     ierr = DMDAVecRestoreArray(dm,F,&f);CHKERRQ(ierr);
178   } break;
179   case ADD_VALUES: {
180     Vec Floc;
181     ierr = DMGetLocalVector(dm,&Floc);CHKERRQ(ierr);
182     ierr = VecZeroEntries(Floc);CHKERRQ(ierr);
183     ierr = DMDAVecGetArray(dm,Floc,&f);CHKERRQ(ierr);
184     CHKMEMQ;
185     ierr = (*dmdats->rhsfunctionlocal)(&info,ptime,x,f,dmdats->rhsfunctionlocalctx);CHKERRQ(ierr);
186     CHKMEMQ;
187     ierr = DMDAVecRestoreArray(dm,Floc,&f);CHKERRQ(ierr);
188     ierr = VecZeroEntries(F);CHKERRQ(ierr);
189     ierr = DMLocalToGlobalBegin(dm,Floc,ADD_VALUES,F);CHKERRQ(ierr);
190     ierr = DMLocalToGlobalEnd(dm,Floc,ADD_VALUES,F);CHKERRQ(ierr);
191     ierr = DMRestoreLocalVector(dm,&Floc);CHKERRQ(ierr);
192   } break;
193   default: SETERRQ1(((PetscObject)ts)->comm,PETSC_ERR_ARG_INCOMP,"Cannot use imode=%d",(int)dmdats->rhsfunctionlocalimode);
194   }
195   ierr = DMDAVecRestoreArray(dm,Xloc,&x);CHKERRQ(ierr);
196   ierr = DMRestoreLocalVector(dm,&Xloc);CHKERRQ(ierr);
197   PetscFunctionReturn(0);
198 }
199 
200 #undef __FUNCT__
201 #define __FUNCT__ "TSComputeRHSJacobian_DMDA"
202 /*
203   This function should eventually replace:
204     DMComputeJacobian() and DMDAComputeJacobian1()
205  */
206 static PetscErrorCode TSComputeRHSJacobian_DMDA(TS ts,PetscReal ptime,Vec X,Mat *A,Mat *B,MatStructure *mstr,void *ctx)
207 {
208   PetscErrorCode ierr;
209   DM             dm;
210   DM_DA_TS     *dmdats = (DM_DA_TS*)ctx;
211   DMDALocalInfo  info;
212   Vec            Xloc;
213   void           *x;
214 
215   PetscFunctionBegin;
216   if (!dmdats->rhsfunctionlocal) SETERRQ(((PetscObject)ts)->comm,PETSC_ERR_PLIB,"Corrupt context");
217   ierr = TSGetDM(ts,&dm);CHKERRQ(ierr);
218 
219   if (dmdats->rhsjacobianlocal) {
220     ierr = DMGetLocalVector(dm,&Xloc);CHKERRQ(ierr);
221     ierr = DMGlobalToLocalBegin(dm,X,INSERT_VALUES,Xloc);CHKERRQ(ierr);
222     ierr = DMGlobalToLocalEnd(dm,X,INSERT_VALUES,Xloc);CHKERRQ(ierr);
223     ierr = DMDAGetLocalInfo(dm,&info);CHKERRQ(ierr);
224     ierr = DMDAVecGetArray(dm,Xloc,&x);CHKERRQ(ierr);
225     CHKMEMQ;
226     ierr = (*dmdats->rhsjacobianlocal)(&info,ptime,x,*A,*B,mstr,dmdats->rhsjacobianlocalctx);CHKERRQ(ierr);
227     CHKMEMQ;
228     ierr = DMDAVecRestoreArray(dm,Xloc,&x);CHKERRQ(ierr);
229     ierr = DMRestoreLocalVector(dm,&Xloc);CHKERRQ(ierr);
230   } else {
231     SETERRQ(((PetscObject)ts)->comm,PETSC_ERR_PLIB,"TSComputeRHSJacobian_DMDA() called without calling DMDATSSetRHSJacobian()");
232   }
233   /* This will be redundant if the user called both, but it's too common to forget. */
234   if (*A != *B) {
235     ierr = MatAssemblyBegin(*A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
236     ierr = MatAssemblyEnd(*A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
237   }
238   PetscFunctionReturn(0);
239 }
240 
241 
242 #undef __FUNCT__
243 #define __FUNCT__ "DMDATSSetRHSFunctionLocal"
244 /*@C
245    DMDATSSetRHSFunctionLocal - set a local residual evaluation function
246 
247    Logically Collective
248 
249    Input Arguments:
250 +  dm - DM to associate callback with
251 .  imode - insert mode for the residual
252 .  func - local residual evaluation
253 -  ctx - optional context for local residual evaluation
254 
255    Calling sequence for func:
256 
257 $ func(DMDALocalInfo info,PetscReal t,void *x,void *f,void *ctx)
258 
259 +  info - DMDALocalInfo defining the subdomain to evaluate the residual on
260 .  t - time at which to evaluate residual
261 .  x - array of local state information
262 .  f - output array of local residual information
263 -  ctx - optional user context
264 
265    Level: beginner
266 
267 .seealso: DMTSSetRHSFunction(), DMDATSSetRHSJacobianLocal(), DMDASNESSetFunctionLocal()
268 @*/
269 PetscErrorCode DMDATSSetRHSFunctionLocal(DM dm,InsertMode imode,DMDATSRHSFunctionLocal func,void *ctx)
270 {
271   PetscErrorCode ierr;
272   TSDM         sdm;
273   DM_DA_TS     *dmdats;
274 
275   PetscFunctionBegin;
276   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
277   ierr = DMTSGetContext(dm,&sdm);CHKERRQ(ierr);
278   ierr = DMDATSGetContext(dm,sdm,&dmdats);CHKERRQ(ierr);
279   dmdats->rhsfunctionlocalimode = imode;
280   dmdats->rhsfunctionlocal = func;
281   dmdats->rhsfunctionlocalctx = ctx;
282   ierr = DMTSSetRHSFunction(dm,TSComputeRHSFunction_DMDA,dmdats);CHKERRQ(ierr);
283   PetscFunctionReturn(0);
284 }
285 
286 #undef __FUNCT__
287 #define __FUNCT__ "DMDATSSetRHSJacobianLocal"
288 /*@C
289    DMDATSSetRHSJacobianLocal - set a local residual evaluation function
290 
291    Logically Collective
292 
293    Input Arguments:
294 +  dm    - DM to associate callback with
295 .  func  - local RHS Jacobian evaluation routine
296 -  ctx   - optional context for local jacobian evaluation
297 
298    Calling sequence for func:
299 
300 $ func(DMDALocalInfo* info,PetscReal t,void* x,Mat J,Mat B,MatStructure *flg,void *ctx);
301 
302 +  info - DMDALocalInfo defining the subdomain to evaluate the residual on
303 .  t    - time at which to evaluate residual
304 .  x    - array of local state information
305 .  J    - Jacobian matrix
306 .  B    - preconditioner matrix; often same as J
307 .  flg  - flag indicating information about the preconditioner matrix structure (same as flag in KSPSetOperators())
308 -  ctx  - optional context passed above
309 
310    Level: beginner
311 
312 .seealso: DMTSSetRHSJacobian(), DMDATSSetRHSFunctionLocal(), DMDASNESSetJacobianLocal()
313 @*/
314 PetscErrorCode DMDATSSetRHSJacobianLocal(DM dm,DMDATSRHSJacobianLocal func,void *ctx)
315 {
316   PetscErrorCode ierr;
317   TSDM         sdm;
318   DM_DA_TS     *dmdats;
319 
320   PetscFunctionBegin;
321   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
322   ierr = DMTSGetContext(dm,&sdm);CHKERRQ(ierr);
323   ierr = DMDATSGetContext(dm,sdm,&dmdats);CHKERRQ(ierr);
324   dmdats->rhsjacobianlocal = func;
325   dmdats->rhsjacobianlocalctx = ctx;
326   ierr = DMTSSetRHSJacobian(dm,TSComputeRHSJacobian_DMDA,dmdats);CHKERRQ(ierr);
327   PetscFunctionReturn(0);
328 }
329 
330 
331 #undef __FUNCT__
332 #define __FUNCT__ "DMDATSSetIFunctionLocal"
333 /*@C
334    DMDATSSetIFunctionLocal - set a local residual evaluation function
335 
336    Logically Collective
337 
338    Input Arguments:
339 +  dm   - DM to associate callback with
340 .  func - local residual evaluation
341 -  ctx  - optional context for local residual evaluation
342 
343    Calling sequence for func:
344 +  info - DMDALocalInfo defining the subdomain to evaluate the residual on
345 .  t    - time at which to evaluate residual
346 .  x    - array of local state information
347 .  xdot - array of local time derivative information
348 .  f    - output array of local function evaluation information
349 -  ctx - optional context passed above
350 
351    Level: beginner
352 
353 .seealso: DMTSSetIFunction(), DMDATSSetIJacobianLocal(), DMDASNESSetFunctionLocal()
354 @*/
355 PetscErrorCode DMDATSSetIFunctionLocal(DM dm,InsertMode imode,DMDATSIFunctionLocal func,void *ctx)
356 {
357   PetscErrorCode ierr;
358   TSDM         sdm;
359   DM_DA_TS     *dmdats;
360 
361   PetscFunctionBegin;
362   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
363   ierr = DMTSGetContext(dm,&sdm);CHKERRQ(ierr);
364   ierr = DMDATSGetContext(dm,sdm,&dmdats);CHKERRQ(ierr);
365   dmdats->ifunctionlocalimode = imode;
366   dmdats->ifunctionlocal = func;
367   dmdats->ifunctionlocalctx = ctx;
368   ierr = DMTSSetIFunction(dm,TSComputeIFunction_DMDA,dmdats);CHKERRQ(ierr);
369   PetscFunctionReturn(0);
370 }
371 
372 #undef __FUNCT__
373 #define __FUNCT__ "DMDATSSetIJacobianLocal"
374 /*@C
375    DMDATSSetIJacobianLocal - set a local residual evaluation function
376 
377    Logically Collective
378 
379    Input Arguments:
380 +  dm   - DM to associate callback with
381 .  func - local residual evaluation
382 -  ctx   - optional context for local residual evaluation
383 
384    Calling sequence for func:
385 
386 $ func(DMDALocalInfo* info,PetscReal t,void* x,void *xdot,Mat J,Mat B,MatStructure *flg,void *ctx);
387 
388 +  info - DMDALocalInfo defining the subdomain to evaluate the residual on
389 .  t    - time at which to evaluate the jacobian
390 .  x    - array of local state information
391 .  xdot - time derivative at this state
392 .  J    - Jacobian matrix
393 .  B    - preconditioner matrix; often same as J
394 .  flg  - flag indicating information about the preconditioner matrix structure (same as flag in KSPSetOperators())
395 -  ctx  - optional context passed above
396 
397    Level: beginner
398 
399 .seealso: DMTSSetJacobian(), DMDATSSetIFunctionLocal(), DMDASNESSetJacobianLocal()
400 @*/
401 PetscErrorCode DMDATSSetIJacobianLocal(DM dm,DMDATSIJacobianLocal func,void *ctx)
402 {
403   PetscErrorCode ierr;
404   TSDM         sdm;
405   DM_DA_TS     *dmdats;
406 
407   PetscFunctionBegin;
408   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
409   ierr = DMTSGetContext(dm,&sdm);CHKERRQ(ierr);
410   ierr = DMDATSGetContext(dm,sdm,&dmdats);CHKERRQ(ierr);
411   dmdats->ijacobianlocal = func;
412   dmdats->ijacobianlocalctx = ctx;
413   ierr = DMTSSetIJacobian(dm,TSComputeIJacobian_DMDA,dmdats);CHKERRQ(ierr);
414   PetscFunctionReturn(0);
415 }
416