xref: /petsc/src/snes/utils/dmdasnes.c (revision 6a5217c03994f2d95bb2e6dbd8bed42381aeb015)
1 #include <petscdmda.h>          /*I "petscdmda.h" I*/
2 #include <petsc/private/dmimpl.h>
3 #include <petsc/private/snesimpl.h>   /*I "petscsnes.h" I*/
4 
5 /* This structure holds the user-provided DMDA callbacks */
6 typedef struct {
7   PetscErrorCode (*residuallocal)(DMDALocalInfo*,void*,void*,void*);
8   PetscErrorCode (*jacobianlocal)(DMDALocalInfo*,void*,Mat,Mat,void*);
9   PetscErrorCode (*objectivelocal)(DMDALocalInfo*,void*,PetscReal*,void*);
10   void       *residuallocalctx;
11   void       *jacobianlocalctx;
12   void       *objectivelocalctx;
13   InsertMode residuallocalimode;
14 
15   /*   For Picard iteration defined locally */
16   PetscErrorCode (*rhsplocal)(DMDALocalInfo*,void*,void*,void*);
17   PetscErrorCode (*jacobianplocal)(DMDALocalInfo*,void*,Mat,Mat,void*);
18   void *picardlocalctx;
19 } DMSNES_DA;
20 
21 static PetscErrorCode DMSNESDestroy_DMDA(DMSNES sdm)
22 {
23   PetscFunctionBegin;
24   PetscCall(PetscFree(sdm->data));
25   PetscFunctionReturn(0);
26 }
27 
28 static PetscErrorCode DMSNESDuplicate_DMDA(DMSNES oldsdm,DMSNES sdm)
29 {
30   PetscFunctionBegin;
31   PetscCall(PetscNewLog(sdm,(DMSNES_DA**)&sdm->data));
32   if (oldsdm->data) {
33     PetscCall(PetscMemcpy(sdm->data,oldsdm->data,sizeof(DMSNES_DA)));
34   }
35   PetscFunctionReturn(0);
36 }
37 
38 static PetscErrorCode DMDASNESGetContext(DM dm,DMSNES sdm,DMSNES_DA  **dmdasnes)
39 {
40   PetscFunctionBegin;
41   *dmdasnes = NULL;
42   if (!sdm->data) {
43     PetscCall(PetscNewLog(dm,(DMSNES_DA**)&sdm->data));
44     sdm->ops->destroy   = DMSNESDestroy_DMDA;
45     sdm->ops->duplicate = DMSNESDuplicate_DMDA;
46   }
47   *dmdasnes = (DMSNES_DA*)sdm->data;
48   PetscFunctionReturn(0);
49 }
50 
51 static PetscErrorCode SNESComputeFunction_DMDA(SNES snes,Vec X,Vec F,void *ctx)
52 {
53   DM             dm;
54   DMSNES_DA      *dmdasnes = (DMSNES_DA*)ctx;
55   DMDALocalInfo  info;
56   Vec            Xloc;
57   void           *x,*f;
58 
59   PetscFunctionBegin;
60   PetscValidHeaderSpecific(snes,SNES_CLASSID,1);
61   PetscValidHeaderSpecific(X,VEC_CLASSID,2);
62   PetscValidHeaderSpecific(F,VEC_CLASSID,3);
63   PetscCheck(dmdasnes->residuallocal,PetscObjectComm((PetscObject)snes),PETSC_ERR_PLIB,"Corrupt context");
64   PetscCall(SNESGetDM(snes,&dm));
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   switch (dmdasnes->residuallocalimode) {
71   case INSERT_VALUES: {
72     PetscCall(DMDAVecGetArray(dm,F,&f));
73     PetscCall(PetscLogEventBegin(SNES_FunctionEval,snes,X,F,0));
74     CHKMEMQ;
75     PetscCall((*dmdasnes->residuallocal)(&info,x,f,dmdasnes->residuallocalctx));
76     CHKMEMQ;
77     PetscCall(PetscLogEventEnd(SNES_FunctionEval,snes,X,F,0));
78     PetscCall(DMDAVecRestoreArray(dm,F,&f));
79   } break;
80   case ADD_VALUES: {
81     Vec Floc;
82     PetscCall(DMGetLocalVector(dm,&Floc));
83     PetscCall(VecZeroEntries(Floc));
84     PetscCall(DMDAVecGetArray(dm,Floc,&f));
85     PetscCall(PetscLogEventBegin(SNES_FunctionEval,snes,X,F,0));
86     CHKMEMQ;
87     PetscCall((*dmdasnes->residuallocal)(&info,x,f,dmdasnes->residuallocalctx));
88     CHKMEMQ;
89     PetscCall(PetscLogEventEnd(SNES_FunctionEval,snes,X,F,0));
90     PetscCall(DMDAVecRestoreArray(dm,Floc,&f));
91     PetscCall(VecZeroEntries(F));
92     PetscCall(DMLocalToGlobalBegin(dm,Floc,ADD_VALUES,F));
93     PetscCall(DMLocalToGlobalEnd(dm,Floc,ADD_VALUES,F));
94     PetscCall(DMRestoreLocalVector(dm,&Floc));
95   } break;
96   default: SETERRQ(PetscObjectComm((PetscObject)snes),PETSC_ERR_ARG_INCOMP,"Cannot use imode=%d",(int)dmdasnes->residuallocalimode);
97   }
98   PetscCall(DMDAVecRestoreArray(dm,Xloc,&x));
99   PetscCall(DMRestoreLocalVector(dm,&Xloc));
100   if (snes->domainerror) {
101     PetscCall(VecSetInf(F));
102   }
103   PetscFunctionReturn(0);
104 }
105 
106 static PetscErrorCode SNESComputeObjective_DMDA(SNES snes,Vec X,PetscReal *ob,void *ctx)
107 {
108   DM             dm;
109   DMSNES_DA      *dmdasnes = (DMSNES_DA*)ctx;
110   DMDALocalInfo  info;
111   Vec            Xloc;
112   void           *x;
113 
114   PetscFunctionBegin;
115   PetscValidHeaderSpecific(snes,SNES_CLASSID,1);
116   PetscValidHeaderSpecific(X,VEC_CLASSID,2);
117   PetscValidRealPointer(ob,3);
118   PetscCheck(dmdasnes->objectivelocal,PetscObjectComm((PetscObject)snes),PETSC_ERR_PLIB,"Corrupt context");
119   PetscCall(SNESGetDM(snes,&dm));
120   PetscCall(DMGetLocalVector(dm,&Xloc));
121   PetscCall(DMGlobalToLocalBegin(dm,X,INSERT_VALUES,Xloc));
122   PetscCall(DMGlobalToLocalEnd(dm,X,INSERT_VALUES,Xloc));
123   PetscCall(DMDAGetLocalInfo(dm,&info));
124   PetscCall(DMDAVecGetArray(dm,Xloc,&x));
125   CHKMEMQ;
126   PetscCall((*dmdasnes->objectivelocal)(&info,x,ob,dmdasnes->objectivelocalctx));
127   CHKMEMQ;
128   PetscCall(DMDAVecRestoreArray(dm,Xloc,&x));
129   PetscCall(DMRestoreLocalVector(dm,&Xloc));
130   PetscFunctionReturn(0);
131 }
132 
133 /* Routine is called by example, hence must be labeled PETSC_EXTERN */
134 PETSC_EXTERN PetscErrorCode SNESComputeJacobian_DMDA(SNES snes,Vec X,Mat A,Mat B,void *ctx)
135 {
136   DM             dm;
137   DMSNES_DA      *dmdasnes = (DMSNES_DA*)ctx;
138   DMDALocalInfo  info;
139   Vec            Xloc;
140   void           *x;
141 
142   PetscFunctionBegin;
143   PetscCheck(dmdasnes->residuallocal,PetscObjectComm((PetscObject)snes),PETSC_ERR_PLIB,"Corrupt context");
144   PetscCall(SNESGetDM(snes,&dm));
145 
146   if (dmdasnes->jacobianlocal) {
147     PetscCall(DMGetLocalVector(dm,&Xloc));
148     PetscCall(DMGlobalToLocalBegin(dm,X,INSERT_VALUES,Xloc));
149     PetscCall(DMGlobalToLocalEnd(dm,X,INSERT_VALUES,Xloc));
150     PetscCall(DMDAGetLocalInfo(dm,&info));
151     PetscCall(DMDAVecGetArray(dm,Xloc,&x));
152     CHKMEMQ;
153     PetscCall((*dmdasnes->jacobianlocal)(&info,x,A,B,dmdasnes->jacobianlocalctx));
154     CHKMEMQ;
155     PetscCall(DMDAVecRestoreArray(dm,Xloc,&x));
156     PetscCall(DMRestoreLocalVector(dm,&Xloc));
157   } else {
158     MatFDColoring fdcoloring;
159     PetscCall(PetscObjectQuery((PetscObject)dm,"DMDASNES_FDCOLORING",(PetscObject*)&fdcoloring));
160     if (!fdcoloring) {
161       ISColoring coloring;
162 
163       PetscCall(DMCreateColoring(dm,dm->coloringtype,&coloring));
164       PetscCall(MatFDColoringCreate(B,coloring,&fdcoloring));
165       switch (dm->coloringtype) {
166       case IS_COLORING_GLOBAL:
167         PetscCall(MatFDColoringSetFunction(fdcoloring,(PetscErrorCode (*)(void))SNESComputeFunction_DMDA,dmdasnes));
168         break;
169       default: SETERRQ(PetscObjectComm((PetscObject)snes),PETSC_ERR_SUP,"No support for coloring type '%s'",ISColoringTypes[dm->coloringtype]);
170       }
171       PetscCall(PetscObjectSetOptionsPrefix((PetscObject)fdcoloring,((PetscObject)dm)->prefix));
172       PetscCall(MatFDColoringSetFromOptions(fdcoloring));
173       PetscCall(MatFDColoringSetUp(B,coloring,fdcoloring));
174       PetscCall(ISColoringDestroy(&coloring));
175       PetscCall(PetscObjectCompose((PetscObject)dm,"DMDASNES_FDCOLORING",(PetscObject)fdcoloring));
176       PetscCall(PetscObjectDereference((PetscObject)fdcoloring));
177 
178       /* The following breaks an ugly reference counting loop that deserves a paragraph. MatFDColoringApply() will call
179        * VecDuplicate() with the state Vec and store inside the MatFDColoring. This Vec will duplicate the Vec, but the
180        * MatFDColoring is composed with the DM. We dereference the DM here so that the reference count will eventually
181        * drop to 0. Note the code in DMDestroy() that exits early for a negative reference count. That code path will be
182        * taken when the PetscObjectList for the Vec inside MatFDColoring is destroyed.
183        */
184       PetscCall(PetscObjectDereference((PetscObject)dm));
185     }
186     PetscCall(MatFDColoringApply(B,fdcoloring,X,snes));
187   }
188   /* This will be redundant if the user called both, but it's too common to forget. */
189   if (A != B) {
190     PetscCall(MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY));
191     PetscCall(MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY));
192   }
193   PetscFunctionReturn(0);
194 }
195 
196 /*@C
197    DMDASNESSetFunctionLocal - set a local residual evaluation function
198 
199    Logically Collective
200 
201    Input Parameters:
202 +  dm - DM to associate callback with
203 .  imode - INSERT_VALUES if local function computes owned part, ADD_VALUES if it contributes to ghosted part
204 .  func - local residual evaluation
205 -  ctx - optional context for local residual evaluation
206 
207    Calling sequence:
208    For PetscErrorCode (*func)(DMDALocalInfo *info,void *x, void *f, void *ctx),
209 +  info - DMDALocalInfo defining the subdomain to evaluate the residual on
210 .  x - dimensional pointer to state at which to evaluate residual (e.g. PetscScalar *x or **x or ***x)
211 .  f - dimensional pointer to residual, write the residual here (e.g. PetscScalar *f or **f or ***f)
212 -  ctx - optional context passed above
213 
214    Level: beginner
215 
216 .seealso: DMDASNESSetJacobianLocal(), DMSNESSetFunction(), DMDACreate1d(), DMDACreate2d(), DMDACreate3d()
217 @*/
218 PetscErrorCode DMDASNESSetFunctionLocal(DM dm,InsertMode imode,PetscErrorCode (*func)(DMDALocalInfo*,void*,void*,void*),void *ctx)
219 {
220   DMSNES         sdm;
221   DMSNES_DA      *dmdasnes;
222 
223   PetscFunctionBegin;
224   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
225   PetscCall(DMGetDMSNESWrite(dm,&sdm));
226   PetscCall(DMDASNESGetContext(dm,sdm,&dmdasnes));
227 
228   dmdasnes->residuallocalimode = imode;
229   dmdasnes->residuallocal      = func;
230   dmdasnes->residuallocalctx   = ctx;
231 
232   PetscCall(DMSNESSetFunction(dm,SNESComputeFunction_DMDA,dmdasnes));
233   if (!sdm->ops->computejacobian) {  /* Call us for the Jacobian too, can be overridden by the user. */
234     PetscCall(DMSNESSetJacobian(dm,SNESComputeJacobian_DMDA,dmdasnes));
235   }
236   PetscFunctionReturn(0);
237 }
238 
239 /*@C
240    DMDASNESSetJacobianLocal - set a local Jacobian evaluation function
241 
242    Logically Collective
243 
244    Input Parameters:
245 +  dm - DM to associate callback with
246 .  func - local Jacobian evaluation
247 -  ctx - optional context for local Jacobian evaluation
248 
249    Calling sequence:
250    For PetscErrorCode (*func)(DMDALocalInfo *info,void *x,Mat J,Mat M,void *ctx),
251 +  info - DMDALocalInfo defining the subdomain to evaluate the Jacobian at
252 .  x - dimensional pointer to state at which to evaluate Jacobian (e.g. PetscScalar *x or **x or ***x)
253 .  J - Mat object for the Jacobian
254 .  M - Mat object for the Jacobian preconditioner matrix
255 -  ctx - optional context passed above
256 
257    Level: beginner
258 
259 .seealso: DMDASNESSetFunctionLocal(), DMSNESSetJacobian(), DMDACreate1d(), DMDACreate2d(), DMDACreate3d()
260 @*/
261 PetscErrorCode DMDASNESSetJacobianLocal(DM dm,PetscErrorCode (*func)(DMDALocalInfo*,void*,Mat,Mat,void*),void *ctx)
262 {
263   DMSNES         sdm;
264   DMSNES_DA      *dmdasnes;
265 
266   PetscFunctionBegin;
267   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
268   PetscCall(DMGetDMSNESWrite(dm,&sdm));
269   PetscCall(DMDASNESGetContext(dm,sdm,&dmdasnes));
270 
271   dmdasnes->jacobianlocal    = func;
272   dmdasnes->jacobianlocalctx = ctx;
273 
274   PetscCall(DMSNESSetJacobian(dm,SNESComputeJacobian_DMDA,dmdasnes));
275   PetscFunctionReturn(0);
276 }
277 
278 /*@C
279    DMDASNESSetObjectiveLocal - set a local residual evaluation function
280 
281    Logically Collective
282 
283    Input Parameters:
284 +  dm - DM to associate callback with
285 .  func - local objective evaluation
286 -  ctx - optional context for local residual evaluation
287 
288    Calling sequence for func:
289 +  info - DMDALocalInfo defining the subdomain to evaluate the residual on
290 .  x - dimensional pointer to state at which to evaluate residual
291 .  ob - eventual objective value
292 -  ctx - optional context passed above
293 
294    Level: beginner
295 
296 .seealso: DMSNESSetFunction(), DMDASNESSetJacobianLocal(), DMDACreate1d(), DMDACreate2d(), DMDACreate3d()
297 @*/
298 PetscErrorCode DMDASNESSetObjectiveLocal(DM dm,DMDASNESObjective func,void *ctx)
299 {
300   DMSNES         sdm;
301   DMSNES_DA      *dmdasnes;
302 
303   PetscFunctionBegin;
304   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
305   PetscCall(DMGetDMSNESWrite(dm,&sdm));
306   PetscCall(DMDASNESGetContext(dm,sdm,&dmdasnes));
307 
308   dmdasnes->objectivelocal    = func;
309   dmdasnes->objectivelocalctx = ctx;
310 
311   PetscCall(DMSNESSetObjective(dm,SNESComputeObjective_DMDA,dmdasnes));
312   PetscFunctionReturn(0);
313 }
314 
315 static PetscErrorCode SNESComputePicard_DMDA(SNES snes,Vec X,Vec F,void *ctx)
316 {
317   DM             dm;
318   DMSNES_DA      *dmdasnes = (DMSNES_DA*)ctx;
319   DMDALocalInfo  info;
320   Vec            Xloc;
321   void           *x,*f;
322 
323   PetscFunctionBegin;
324   PetscValidHeaderSpecific(snes,SNES_CLASSID,1);
325   PetscValidHeaderSpecific(X,VEC_CLASSID,2);
326   PetscValidHeaderSpecific(F,VEC_CLASSID,3);
327   PetscCheck(dmdasnes->rhsplocal,PetscObjectComm((PetscObject)snes),PETSC_ERR_PLIB,"Corrupt context");
328   PetscCall(SNESGetDM(snes,&dm));
329   PetscCall(DMGetLocalVector(dm,&Xloc));
330   PetscCall(DMGlobalToLocalBegin(dm,X,INSERT_VALUES,Xloc));
331   PetscCall(DMGlobalToLocalEnd(dm,X,INSERT_VALUES,Xloc));
332   PetscCall(DMDAGetLocalInfo(dm,&info));
333   PetscCall(DMDAVecGetArray(dm,Xloc,&x));
334   switch (dmdasnes->residuallocalimode) {
335   case INSERT_VALUES: {
336     PetscCall(DMDAVecGetArray(dm,F,&f));
337     CHKMEMQ;
338     PetscCall((*dmdasnes->rhsplocal)(&info,x,f,dmdasnes->picardlocalctx));
339     CHKMEMQ;
340     PetscCall(DMDAVecRestoreArray(dm,F,&f));
341   } break;
342   case ADD_VALUES: {
343     Vec Floc;
344     PetscCall(DMGetLocalVector(dm,&Floc));
345     PetscCall(VecZeroEntries(Floc));
346     PetscCall(DMDAVecGetArray(dm,Floc,&f));
347     CHKMEMQ;
348     PetscCall((*dmdasnes->rhsplocal)(&info,x,f,dmdasnes->picardlocalctx));
349     CHKMEMQ;
350     PetscCall(DMDAVecRestoreArray(dm,Floc,&f));
351     PetscCall(VecZeroEntries(F));
352     PetscCall(DMLocalToGlobalBegin(dm,Floc,ADD_VALUES,F));
353     PetscCall(DMLocalToGlobalEnd(dm,Floc,ADD_VALUES,F));
354     PetscCall(DMRestoreLocalVector(dm,&Floc));
355   } break;
356   default: SETERRQ(PetscObjectComm((PetscObject)snes),PETSC_ERR_ARG_INCOMP,"Cannot use imode=%d",(int)dmdasnes->residuallocalimode);
357   }
358   PetscCall(DMDAVecRestoreArray(dm,Xloc,&x));
359   PetscCall(DMRestoreLocalVector(dm,&Xloc));
360   PetscFunctionReturn(0);
361 }
362 
363 static PetscErrorCode SNESComputePicardJacobian_DMDA(SNES snes,Vec X,Mat A,Mat B,void *ctx)
364 {
365   DM             dm;
366   DMSNES_DA      *dmdasnes = (DMSNES_DA*)ctx;
367   DMDALocalInfo  info;
368   Vec            Xloc;
369   void           *x;
370 
371   PetscFunctionBegin;
372   PetscCheck(dmdasnes->jacobianplocal,PetscObjectComm((PetscObject)snes),PETSC_ERR_PLIB,"Corrupt context");
373   PetscCall(SNESGetDM(snes,&dm));
374 
375   PetscCall(DMGetLocalVector(dm,&Xloc));
376   PetscCall(DMGlobalToLocalBegin(dm,X,INSERT_VALUES,Xloc));
377   PetscCall(DMGlobalToLocalEnd(dm,X,INSERT_VALUES,Xloc));
378   PetscCall(DMDAGetLocalInfo(dm,&info));
379   PetscCall(DMDAVecGetArray(dm,Xloc,&x));
380   CHKMEMQ;
381   PetscCall((*dmdasnes->jacobianplocal)(&info,x,A,B,dmdasnes->picardlocalctx));
382   CHKMEMQ;
383   PetscCall(DMDAVecRestoreArray(dm,Xloc,&x));
384   PetscCall(DMRestoreLocalVector(dm,&Xloc));
385   if (A != B) {
386     PetscCall(MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY));
387     PetscCall(MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY));
388   }
389   PetscFunctionReturn(0);
390 }
391 
392 /*@C
393    DMDASNESSetPicardLocal - set a local right hand side and matrix evaluation function for Picard iteration
394 
395    Logically Collective
396 
397    Input Parameters:
398 +  dm - DM to associate callback with
399 .  imode - INSERT_VALUES if local function computes owned part, ADD_VALUES if it contributes to ghosted part
400 .  func - local residual evaluation
401 -  ctx - optional context for local residual evaluation
402 
403    Calling sequence for func:
404 +  info - DMDALocalInfo defining the subdomain to evaluate the residual on
405 .  x - dimensional pointer to state at which to evaluate residual
406 .  f - dimensional pointer to residual, write the residual here
407 -  ctx - optional context passed above
408 
409    Notes:
410     The user must use
411     PetscCall(SNESSetFunction(snes,NULL,SNESPicardComputeFunction,&user));
412     in their code before calling this routine.
413 
414    Level: beginner
415 
416 .seealso: DMSNESSetFunction(), DMDASNESSetJacobian(), DMDACreate1d(), DMDACreate2d(), DMDACreate3d()
417 @*/
418 PetscErrorCode DMDASNESSetPicardLocal(DM dm,InsertMode imode,PetscErrorCode (*func)(DMDALocalInfo*,void*,void*,void*),
419                                       PetscErrorCode (*jac)(DMDALocalInfo*,void*,Mat,Mat,void*),void *ctx)
420 {
421   DMSNES         sdm;
422   DMSNES_DA      *dmdasnes;
423 
424   PetscFunctionBegin;
425   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
426   PetscCall(DMGetDMSNESWrite(dm,&sdm));
427   PetscCall(DMDASNESGetContext(dm,sdm,&dmdasnes));
428 
429   dmdasnes->residuallocalimode = imode;
430   dmdasnes->rhsplocal          = func;
431   dmdasnes->jacobianplocal     = jac;
432   dmdasnes->picardlocalctx     = ctx;
433 
434   PetscCall(DMSNESSetPicard(dm,SNESComputePicard_DMDA,SNESComputePicardJacobian_DMDA,dmdasnes));
435   PetscCall(DMSNESSetMFFunction(dm,SNESComputeFunction_DMDA,dmdasnes));
436   PetscFunctionReturn(0);
437 }
438