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