xref: /petsc/src/ts/utils/dmdats.c (revision 0e9bae810fdaeb60e2713eaa8ddb89f42e079fd1)
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 .  func - local residual evaluation
252 -  ctx - optional context for local residual evaluation
253 
254    Calling sequence for func:
255 +  info - DMDALocalInfo defining the subdomain to evaluate the residual on
256 .  imode - INSERT_VALUES if local function computes owned part, ADD_VALUES if it contributes to ghosted part
257 .  x - dimensional pointer to state at which to evaluate residual
258 .  f - dimensional pointer to residual, write the residual here
259 -  ctx - optional context passed above
260 
261    Level: beginner
262 
263 .seealso: DMTSSetFunction(), DMDATSSetJacobian(), DMDACreate1d(), DMDACreate2d(), DMDACreate3d()
264 @*/
265 PetscErrorCode DMDATSSetRHSFunctionLocal(DM dm,InsertMode imode,DMDATSRHSFunctionLocal func,void *ctx)
266 {
267   PetscErrorCode ierr;
268   TSDM         sdm;
269   DM_DA_TS     *dmdats;
270 
271   PetscFunctionBegin;
272   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
273   ierr = DMTSGetContext(dm,&sdm);CHKERRQ(ierr);
274   ierr = DMDATSGetContext(dm,sdm,&dmdats);CHKERRQ(ierr);
275   dmdats->rhsfunctionlocalimode = imode;
276   dmdats->rhsfunctionlocal = func;
277   dmdats->rhsfunctionlocalctx = ctx;
278   ierr = DMTSSetRHSFunction(dm,TSComputeRHSFunction_DMDA,dmdats);CHKERRQ(ierr);
279   PetscFunctionReturn(0);
280 }
281 
282 #undef __FUNCT__
283 #define __FUNCT__ "DMDATSSetRHSJacobianLocal"
284 /*@C
285    DMDATSSetRHSJacobianLocal - set a local residual evaluation function
286 
287    Logically Collective
288 
289    Input Arguments:
290 +  dm - DM to associate callback with
291 .  func - local residual evaluation
292 -  ctx - optional context for local residual evaluation
293 
294    Calling sequence for func:
295 +  info - DMDALocalInfo defining the subdomain to evaluate the residual on
296 .  x - dimensional pointer to state at which to evaluate residual
297 .  f - dimensional pointer to residual, write the residual here
298 -  ctx - optional context passed above
299 
300    Level: beginner
301 
302 .seealso: DMTSSetJacobian(), DMDATSSetJacobian(), DMDACreate1d(), DMDACreate2d(), DMDACreate3d()
303 @*/
304 PetscErrorCode DMDATSSetRHSJacobianLocal(DM dm,DMDATSRHSJacobianLocal func,void *ctx)
305 {
306   PetscErrorCode ierr;
307   TSDM         sdm;
308   DM_DA_TS     *dmdats;
309 
310   PetscFunctionBegin;
311   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
312   ierr = DMTSGetContext(dm,&sdm);CHKERRQ(ierr);
313   ierr = DMDATSGetContext(dm,sdm,&dmdats);CHKERRQ(ierr);
314   dmdats->rhsjacobianlocal = func;
315   dmdats->rhsjacobianlocalctx = ctx;
316   ierr = DMTSSetRHSJacobian(dm,TSComputeRHSJacobian_DMDA,dmdats);CHKERRQ(ierr);
317   PetscFunctionReturn(0);
318 }
319 
320 
321 #undef __FUNCT__
322 #define __FUNCT__ "DMDATSSetIFunctionLocal"
323 /*@C
324    DMDATSSetIFunctionLocal - set a local residual evaluation function
325 
326    Logically Collective
327 
328    Input Arguments:
329 +  dm - DM to associate callback with
330 .  func - local residual evaluation
331 -  ctx - optional context for local residual evaluation
332 
333    Calling sequence for func:
334 +  info - DMDALocalInfo defining the subdomain to evaluate the residual on
335 .  imode - INSERT_VALUES if local function computes owned part, ADD_VALUES if it contributes to ghosted part
336 .  x - dimensional pointer to state at which to evaluate residual
337 .  f - dimensional pointer to residual, write the residual here
338 -  ctx - optional context passed above
339 
340    Level: beginner
341 
342 .seealso: DMTSSetFunction(), DMDATSSetJacobian(), DMDACreate1d(), DMDACreate2d(), DMDACreate3d()
343 @*/
344 PetscErrorCode DMDATSSetIFunctionLocal(DM dm,InsertMode imode,DMDATSIFunctionLocal func,void *ctx)
345 {
346   PetscErrorCode ierr;
347   TSDM         sdm;
348   DM_DA_TS     *dmdats;
349 
350   PetscFunctionBegin;
351   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
352   ierr = DMTSGetContext(dm,&sdm);CHKERRQ(ierr);
353   ierr = DMDATSGetContext(dm,sdm,&dmdats);CHKERRQ(ierr);
354   dmdats->ifunctionlocalimode = imode;
355   dmdats->ifunctionlocal = func;
356   dmdats->ifunctionlocalctx = ctx;
357   ierr = DMTSSetIFunction(dm,TSComputeIFunction_DMDA,dmdats);CHKERRQ(ierr);
358   PetscFunctionReturn(0);
359 }
360 
361 #undef __FUNCT__
362 #define __FUNCT__ "DMDATSSetIJacobianLocal"
363 /*@C
364    DMDATSSetIJacobianLocal - set a local residual evaluation function
365 
366    Logically Collective
367 
368    Input Arguments:
369 +  dm - DM to associate callback with
370 .  func - local residual evaluation
371 -  ctx - optional context for local residual evaluation
372 
373    Calling sequence for func:
374 +  info - DMDALocalInfo defining the subdomain to evaluate the residual on
375 .  x - dimensional pointer to state at which to evaluate residual
376 .  f - dimensional pointer to residual, write the residual here
377 -  ctx - optional context passed above
378 
379    Level: beginner
380 
381 .seealso: DMTSSetJacobian(), DMDATSSetJacobian(), DMDACreate1d(), DMDACreate2d(), DMDACreate3d()
382 @*/
383 PetscErrorCode DMDATSSetIJacobianLocal(DM dm,DMDATSIJacobianLocal func,void *ctx)
384 {
385   PetscErrorCode ierr;
386   TSDM         sdm;
387   DM_DA_TS     *dmdats;
388 
389   PetscFunctionBegin;
390   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
391   ierr = DMTSGetContext(dm,&sdm);CHKERRQ(ierr);
392   ierr = DMDATSGetContext(dm,sdm,&dmdats);CHKERRQ(ierr);
393   dmdats->ijacobianlocal = func;
394   dmdats->ijacobianlocalctx = ctx;
395   ierr = DMTSSetIJacobian(dm,TSComputeIJacobian_DMDA,dmdats);CHKERRQ(ierr);
396   PetscFunctionReturn(0);
397 }
398