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