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