xref: /petsc/src/snes/utils/dmdasnes.c (revision 5edf65166e22df6feda2ed6bd8a6e348419dee17)
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,MatStructure*,void*);
9   void *residuallocalctx;
10   void *jacobianlocalctx;
11   InsertMode residuallocalimode;
12 } DM_DA_SNES;
13 
14 #undef __FUNCT__
15 #define __FUNCT__ "SNESDMDestroy_DMDA"
16 static PetscErrorCode SNESDMDestroy_DMDA(SNESDM sdm)
17 {
18   PetscErrorCode ierr;
19 
20   PetscFunctionBegin;
21   ierr = PetscFree(sdm->data);CHKERRQ(ierr);
22   PetscFunctionReturn(0);
23 }
24 
25 #undef __FUNCT__
26 #define __FUNCT__ "SNESDMDuplicate_DMDA"
27 static PetscErrorCode SNESDMDuplicate_DMDA(SNESDM oldsdm,DM dm)
28 {
29   PetscErrorCode ierr;
30   SNESDM         sdm;
31 
32   PetscFunctionBegin;
33   ierr = DMSNESGetContext(dm,&sdm);CHKERRQ(ierr);
34   ierr = PetscNewLog(dm,DM_DA_SNES,&sdm->data);CHKERRQ(ierr);
35   if (oldsdm->data)ierr = PetscMemcpy(sdm->data,oldsdm->data,sizeof(DM_DA_SNES));CHKERRQ(ierr);
36   PetscFunctionReturn(0);
37 }
38 
39 
40 #undef __FUNCT__
41 #define __FUNCT__ "DMDASNESGetContext"
42 static PetscErrorCode DMDASNESGetContext(DM dm,SNESDM sdm,DM_DA_SNES **dmdasnes)
43 {
44   PetscErrorCode ierr;
45 
46   PetscFunctionBegin;
47   *dmdasnes = PETSC_NULL;
48   if (!sdm->data) {
49     ierr = PetscNewLog(dm,DM_DA_SNES,&sdm->data);CHKERRQ(ierr);
50     sdm->destroy = SNESDMDestroy_DMDA;
51     sdm->duplicate = SNESDMDuplicate_DMDA;
52   }
53   *dmdasnes = (DM_DA_SNES*)sdm->data;
54   PetscFunctionReturn(0);
55 }
56 
57 #undef __FUNCT__
58 #define __FUNCT__ "SNESComputeFunction_DMDA"
59 /*
60   This function should eventually replace:
61     DMDAComputeFunction() and DMDAComputeFunction1()
62  */
63 static PetscErrorCode SNESComputeFunction_DMDA(SNES snes,Vec X,Vec F,void *ctx)
64 {
65   PetscErrorCode ierr;
66   DM             dm;
67   DM_DA_SNES     *dmdasnes = (DM_DA_SNES*)ctx;
68   DMDALocalInfo  info;
69   Vec            Xloc;
70   void           *x,*f;
71 
72   PetscFunctionBegin;
73   PetscValidHeaderSpecific(snes,SNES_CLASSID,1);
74   PetscValidHeaderSpecific(X,VEC_CLASSID,2);
75   PetscValidHeaderSpecific(F,VEC_CLASSID,3);
76   if (!dmdasnes->residuallocal) SETERRQ(((PetscObject)snes)->comm,PETSC_ERR_PLIB,"Corrupt context");
77   ierr = SNESGetDM(snes,&dm);CHKERRQ(ierr);
78   ierr = DMGetLocalVector(dm,&Xloc);CHKERRQ(ierr);
79   ierr = DMGlobalToLocalBegin(dm,X,INSERT_VALUES,Xloc);CHKERRQ(ierr);
80   ierr = DMGlobalToLocalEnd(dm,X,INSERT_VALUES,Xloc);CHKERRQ(ierr);
81   ierr = DMDAGetLocalInfo(dm,&info);CHKERRQ(ierr);
82   ierr = DMDAVecGetArray(dm,Xloc,&x);CHKERRQ(ierr);
83   switch (dmdasnes->residuallocalimode) {
84   case INSERT_VALUES: {
85     ierr = DMDAVecGetArray(dm,F,&f);CHKERRQ(ierr);
86     CHKMEMQ;
87     ierr = (*dmdasnes->residuallocal)(&info,x,f,dmdasnes->residuallocalctx);CHKERRQ(ierr);
88     CHKMEMQ;
89     ierr = DMDAVecRestoreArray(dm,F,&f);CHKERRQ(ierr);
90   } break;
91   case ADD_VALUES: {
92     Vec Floc;
93     ierr = DMGetLocalVector(dm,&Floc);CHKERRQ(ierr);
94     ierr = VecZeroEntries(Floc);CHKERRQ(ierr);
95     ierr = DMDAVecGetArray(dm,Floc,&f);CHKERRQ(ierr);
96     CHKMEMQ;
97     ierr = (*dmdasnes->residuallocal)(&info,x,f,dmdasnes->residuallocalctx);CHKERRQ(ierr);
98     CHKMEMQ;
99     ierr = DMDAVecRestoreArray(dm,Floc,&f);CHKERRQ(ierr);
100     ierr = VecZeroEntries(F);CHKERRQ(ierr);
101     ierr = DMLocalToGlobalBegin(dm,Floc,ADD_VALUES,F);CHKERRQ(ierr);
102     ierr = DMLocalToGlobalEnd(dm,Floc,ADD_VALUES,F);CHKERRQ(ierr);
103     ierr = DMRestoreLocalVector(dm,&Floc);CHKERRQ(ierr);
104   } break;
105   default: SETERRQ1(((PetscObject)snes)->comm,PETSC_ERR_ARG_INCOMP,"Cannot use imode=%d",(int)dmdasnes->residuallocalimode);
106   }
107   ierr = DMDAVecRestoreArray(dm,Xloc,&x);CHKERRQ(ierr);
108   ierr = DMRestoreLocalVector(dm,&Xloc);CHKERRQ(ierr);
109   PetscFunctionReturn(0);
110 }
111 
112 #undef __FUNCT__
113 #define __FUNCT__ "SNESComputeLocalBlockFunction_DMDA"
114 /*
115   This function should eventually replace:
116     DMDAComputeFunction() and DMDAComputeFunction1()
117  */
118 PetscErrorCode SNESComputeLocalBlockFunction_DMDA(SNES snes,Vec X,Vec F,void *dmctx)
119 {
120   PetscErrorCode ierr;
121   DM             dm = (DM)dmctx; /* the global DMDA context */
122   DMDALocalInfo  info;
123   void           *x,*f;
124   DM_DA_SNES     *dmdasnes;
125   SNESDM         sdm;
126 
127   PetscFunctionBegin;
128 
129   PetscValidHeaderSpecific(snes,SNES_CLASSID,1);
130   PetscValidHeaderSpecific(X,VEC_CLASSID,2);
131   PetscValidHeaderSpecific(F,VEC_CLASSID,3);
132   PetscValidHeaderSpecific(dm,DM_CLASSID,4);
133 
134   ierr = DMSNESGetContext(dm,&sdm);CHKERRQ(ierr);
135   ierr = DMDASNESGetContext(dm,sdm,&dmdasnes);CHKERRQ(ierr);
136   if (!dmdasnes->residuallocal) SETERRQ(((PetscObject)snes)->comm,PETSC_ERR_PLIB,"Corrupt context");
137   ierr = DMDAVecGetArray(dm,X,&x);CHKERRQ(ierr);
138   ierr = DMDAGetLocalBlockInfo(dm,&info);CHKERRQ(ierr);
139   switch (dmdasnes->residuallocalimode) {
140   case INSERT_VALUES: {
141     ierr = DMDAVecGetArray(dm,F,&f);CHKERRQ(ierr);
142     CHKMEMQ;
143     ierr = (*dmdasnes->residuallocal)(&info,x,f,dmdasnes->residuallocalctx);CHKERRQ(ierr);
144     CHKMEMQ;
145     ierr = DMDAVecRestoreArray(dm,F,&f);CHKERRQ(ierr);
146   } break;
147   case ADD_VALUES: {
148     Vec Floc;
149     ierr = DMGetLocalVector(dm,&Floc);CHKERRQ(ierr);
150     ierr = VecZeroEntries(Floc);CHKERRQ(ierr);
151     ierr = DMDAVecGetArray(dm,Floc,&f);CHKERRQ(ierr);
152     CHKMEMQ;
153     ierr = (*dmdasnes->residuallocal)(&info,x,f,dmdasnes->residuallocalctx);CHKERRQ(ierr);
154     CHKMEMQ;
155     ierr = DMDAVecRestoreArray(dm,Floc,&f);CHKERRQ(ierr);
156     ierr = VecZeroEntries(F);CHKERRQ(ierr);
157     ierr = DMLocalToGlobalBegin(dm,Floc,ADD_VALUES,F);CHKERRQ(ierr);
158     ierr = DMLocalToGlobalEnd(dm,Floc,ADD_VALUES,F);CHKERRQ(ierr);
159     ierr = DMRestoreLocalVector(dm,&Floc);CHKERRQ(ierr);
160   } break;
161   default: SETERRQ1(((PetscObject)snes)->comm,PETSC_ERR_ARG_INCOMP,"Cannot use imode=%d",(int)dmdasnes->residuallocalimode);
162   }
163   ierr = DMDAVecRestoreArray(dm,X,&x);CHKERRQ(ierr);
164   PetscFunctionReturn(0);
165 }
166 
167 #undef __FUNCT__
168 #define __FUNCT__ "SNESComputeJacobian_DMDA"
169 /*
170   This function should eventually replace:
171     DMComputeJacobian() and DMDAComputeJacobian1()
172  */
173 static PetscErrorCode SNESComputeJacobian_DMDA(SNES snes,Vec X,Mat *A,Mat *B,MatStructure *mstr,void *ctx)
174 {
175   PetscErrorCode ierr;
176   DM             dm;
177   DM_DA_SNES     *dmdasnes = (DM_DA_SNES*)ctx;
178   DMDALocalInfo  info;
179   Vec            Xloc;
180   void           *x;
181 
182   PetscFunctionBegin;
183   if (!dmdasnes->residuallocal) SETERRQ(((PetscObject)snes)->comm,PETSC_ERR_PLIB,"Corrupt context");
184   ierr = SNESGetDM(snes,&dm);CHKERRQ(ierr);
185 
186   if (dmdasnes->jacobianlocal) {
187     ierr = DMGetLocalVector(dm,&Xloc);CHKERRQ(ierr);
188     ierr = DMGlobalToLocalBegin(dm,X,INSERT_VALUES,Xloc);CHKERRQ(ierr);
189     ierr = DMGlobalToLocalEnd(dm,X,INSERT_VALUES,Xloc);CHKERRQ(ierr);
190     ierr = DMDAGetLocalInfo(dm,&info);CHKERRQ(ierr);
191     ierr = DMDAVecGetArray(dm,Xloc,&x);CHKERRQ(ierr);
192     CHKMEMQ;
193     ierr = (*dmdasnes->jacobianlocal)(&info,x,*A,*B,mstr,dmdasnes->jacobianlocalctx);CHKERRQ(ierr);
194     CHKMEMQ;
195     ierr = DMDAVecRestoreArray(dm,Xloc,&x);CHKERRQ(ierr);
196     ierr = DMRestoreLocalVector(dm,&Xloc);CHKERRQ(ierr);
197   } else {
198     MatFDColoring fdcoloring;
199     ierr = PetscObjectQuery((PetscObject)dm,"DMDASNES_FDCOLORING",(PetscObject*)&fdcoloring);CHKERRQ(ierr);
200     if (!fdcoloring) {
201       ISColoring     coloring;
202 
203       ierr = DMCreateColoring(dm,dm->coloringtype,MATAIJ,&coloring);CHKERRQ(ierr);
204       ierr = MatFDColoringCreate(*B,coloring,&fdcoloring);CHKERRQ(ierr);
205       ierr = ISColoringDestroy(&coloring);CHKERRQ(ierr);
206       switch (dm->coloringtype) {
207       case IS_COLORING_GLOBAL:
208         ierr = MatFDColoringSetFunction(fdcoloring,(PetscErrorCode(*)(void))SNESComputeFunction_DMDA,dmdasnes);CHKERRQ(ierr);
209         break;
210       default: SETERRQ1(((PetscObject)snes)->comm,PETSC_ERR_SUP,"No support for coloring type '%s'",ISColoringTypes[dm->coloringtype]);
211       }
212       ierr = PetscObjectSetOptionsPrefix((PetscObject)fdcoloring,((PetscObject)dm)->prefix);CHKERRQ(ierr);
213       ierr = MatFDColoringSetFromOptions(fdcoloring);CHKERRQ(ierr);
214       ierr = PetscObjectCompose((PetscObject)dm,"DMDASNES_FDCOLORING",(PetscObject)fdcoloring);CHKERRQ(ierr);
215       ierr = PetscObjectDereference((PetscObject)fdcoloring);CHKERRQ(ierr);
216 
217       /* The following breaks an ugly reference counting loop that deserves a paragraph. MatFDColoringApply() will call
218        * VecDuplicate() with the state Vec and store inside the MatFDColoring. This Vec will duplicate the Vec, but the
219        * MatFDColoring is composed with the DM. We dereference the DM here so that the reference count will eventually
220        * drop to 0. Note the code in DMDestroy() that exits early for a negative reference count. That code path will be
221        * taken when the PetscOList for the Vec inside MatFDColoring is destroyed.
222        */
223       ierr = PetscObjectDereference((PetscObject)dm);CHKERRQ(ierr);
224     }
225     *mstr = SAME_NONZERO_PATTERN;
226     ierr = MatFDColoringApply(*B,fdcoloring,X,mstr,snes);CHKERRQ(ierr);
227   }
228   /* This will be redundant if the user called both, but it's too common to forget. */
229   if (*A != *B) {
230     ierr = MatAssemblyBegin(*A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
231     ierr = MatAssemblyEnd(*A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
232   }
233   PetscFunctionReturn(0);
234 }
235 
236 #undef __FUNCT__
237 #define __FUNCT__ "SNESComputeJacobian_DMDA"
238 /*
239   This function should eventually replace:
240     DMComputeJacobian() and DMDAComputeJacobian1()
241  */
242 PetscErrorCode SNESComputeLocalBlockJacobian_DMDA(SNES snes,Vec X,Mat *A,Mat *B,MatStructure *mstr,void *dmctx)
243 {
244   PetscErrorCode ierr;
245   DM             dm = (DM)dmctx; /* the global DMDA context */
246   DM_DA_SNES     *dmdasnes;
247   SNESDM         sdm;
248   DMDALocalInfo  info;
249   void           *x;
250 
251   PetscFunctionBegin;
252   ierr = DMSNESGetContext(dm,&sdm);CHKERRQ(ierr);
253   ierr = DMDASNESGetContext(dm,sdm,&dmdasnes);CHKERRQ(ierr);
254   if (!dmdasnes->residuallocal) SETERRQ(((PetscObject)snes)->comm,PETSC_ERR_PLIB,"Corrupt context");
255   if (dmdasnes->jacobianlocal) {
256     ierr = DMDAVecGetArray(dm,X,&x);CHKERRQ(ierr);
257     CHKMEMQ;
258     ierr = (*dmdasnes->jacobianlocal)(&info,x,*A,*B,mstr,dmdasnes->jacobianlocalctx);CHKERRQ(ierr);
259     CHKMEMQ;
260     ierr = DMDAVecRestoreArray(dm,X,&x);CHKERRQ(ierr);
261   }
262   /* This will be redundant if the user called both, but it's too common to forget. */
263   if (*A != *B) {
264     ierr = MatAssemblyBegin(*A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
265     ierr = MatAssemblyEnd(*A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
266   }
267   PetscFunctionReturn(0);
268 }
269 
270 
271 #undef __FUNCT__
272 #define __FUNCT__ "DMDASNESSetFunctionLocal"
273 /*@C
274    DMDASNESSetFunctionLocal - set a local residual evaluation function
275 
276    Logically Collective
277 
278    Input Arguments:
279 +  dm - DM to associate callback with
280 .  func - local residual evaluation
281 -  ctx - optional context for local residual evaluation
282 
283    Calling sequence for func:
284 +  info - DMDALocalInfo defining the subdomain to evaluate the residual on
285 .  imode - INSERT_VALUES if local function computes owned part, ADD_VALUES if it contributes to ghosted part
286 .  x - dimensional pointer to state at which to evaluate residual
287 .  f - dimensional pointer to residual, write the residual here
288 -  ctx - optional context passed above
289 
290    Level: beginner
291 
292 .seealso: DMSNESSetFunction(), DMDASNESSetJacobian(), DMDACreate1d(), DMDACreate2d(), DMDACreate3d()
293 @*/
294 PetscErrorCode DMDASNESSetFunctionLocal(DM dm,InsertMode imode,PetscErrorCode (*func)(DMDALocalInfo*,void*,void*,void*),void *ctx)
295 {
296   PetscErrorCode ierr;
297   SNESDM         sdm;
298   DM_DA_SNES     *dmdasnes;
299 
300   PetscFunctionBegin;
301   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
302   ierr = DMSNESGetContext(dm,&sdm);CHKERRQ(ierr);
303   ierr = DMDASNESGetContext(dm,sdm,&dmdasnes);CHKERRQ(ierr);
304   dmdasnes->residuallocalimode = imode;
305   dmdasnes->residuallocal = func;
306   dmdasnes->residuallocalctx = ctx;
307   ierr = DMSNESSetFunction(dm,SNESComputeFunction_DMDA,dmdasnes);CHKERRQ(ierr);
308   ierr = DMSNESSetBlockFunction(dm,SNESComputeLocalBlockFunction_DMDA,dm);CHKERRQ(ierr);
309   if (!sdm->computejacobian) {  /* Call us for the Jacobian too, can be overridden by the user. */
310     ierr = DMSNESSetJacobian(dm,SNESComputeJacobian_DMDA,dmdasnes);CHKERRQ(ierr);
311   }
312   PetscFunctionReturn(0);
313 }
314 
315 #undef __FUNCT__
316 #define __FUNCT__ "DMDASNESSetJacobianLocal"
317 /*@C
318    DMDASNESSetJacobianLocal - set a local residual evaluation function
319 
320    Logically Collective
321 
322    Input Arguments:
323 +  dm - DM to associate callback with
324 .  func - local residual evaluation
325 -  ctx - optional context for local residual evaluation
326 
327    Calling sequence for func:
328 +  info - DMDALocalInfo defining the subdomain to evaluate the residual on
329 .  x - dimensional pointer to state at which to evaluate residual
330 .  f - dimensional pointer to residual, write the residual here
331 -  ctx - optional context passed above
332 
333    Level: beginner
334 
335 .seealso: DMSNESSetJacobian(), DMDASNESSetJacobian(), DMDACreate1d(), DMDACreate2d(), DMDACreate3d()
336 @*/
337 PetscErrorCode DMDASNESSetJacobianLocal(DM dm,PetscErrorCode (*func)(DMDALocalInfo*,void*,Mat,Mat,MatStructure*,void*),void *ctx)
338 {
339   PetscErrorCode ierr;
340   SNESDM         sdm;
341   DM_DA_SNES     *dmdasnes;
342 
343   PetscFunctionBegin;
344   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
345   ierr = DMSNESGetContext(dm,&sdm);CHKERRQ(ierr);
346   ierr = DMDASNESGetContext(dm,sdm,&dmdasnes);CHKERRQ(ierr);
347   dmdasnes->jacobianlocal = func;
348   dmdasnes->jacobianlocalctx = ctx;
349   ierr = DMSNESSetJacobian(dm,SNESComputeJacobian_DMDA,dmdasnes);CHKERRQ(ierr);
350   ierr = DMSNESSetBlockJacobian(dm,SNESComputeLocalBlockJacobian_DMDA,dm);CHKERRQ(ierr);
351   PetscFunctionReturn(0);
352 }
353