xref: /petsc/src/snes/utils/dmlocalsnes.c (revision df4cd43f92eaa320656440c40edb1046daee8f75)
1 #include <petsc/private/dmimpl.h>
2 #include <petsc/private/snesimpl.h> /*I "petscsnes.h" I*/
3 
4 typedef struct {
5   PetscErrorCode (*residuallocal)(DM, Vec, Vec, void *);
6   PetscErrorCode (*jacobianlocal)(DM, Vec, Mat, Mat, void *);
7   PetscErrorCode (*boundarylocal)(DM, Vec, void *);
8   void *residuallocalctx;
9   void *jacobianlocalctx;
10   void *boundarylocalctx;
11 } DMSNES_Local;
12 
13 static PetscErrorCode DMSNESDestroy_DMLocal(DMSNES sdm)
14 {
15   PetscFunctionBegin;
16   PetscCall(PetscFree(sdm->data));
17   sdm->data = NULL;
18   PetscFunctionReturn(PETSC_SUCCESS);
19 }
20 
21 static PetscErrorCode DMSNESDuplicate_DMLocal(DMSNES oldsdm, DMSNES sdm)
22 {
23   PetscFunctionBegin;
24   if (sdm->data != oldsdm->data) {
25     PetscCall(PetscFree(sdm->data));
26     PetscCall(PetscNew((DMSNES_Local **)&sdm->data));
27     if (oldsdm->data) PetscCall(PetscMemcpy(sdm->data, oldsdm->data, sizeof(DMSNES_Local)));
28   }
29   PetscFunctionReturn(PETSC_SUCCESS);
30 }
31 
32 static PetscErrorCode DMLocalSNESGetContext(DM dm, DMSNES sdm, DMSNES_Local **dmlocalsnes)
33 {
34   PetscFunctionBegin;
35   *dmlocalsnes = NULL;
36   if (!sdm->data) {
37     PetscCall(PetscNew((DMSNES_Local **)&sdm->data));
38 
39     sdm->ops->destroy   = DMSNESDestroy_DMLocal;
40     sdm->ops->duplicate = DMSNESDuplicate_DMLocal;
41   }
42   *dmlocalsnes = (DMSNES_Local *)sdm->data;
43   PetscFunctionReturn(PETSC_SUCCESS);
44 }
45 
46 static PetscErrorCode SNESComputeFunction_DMLocal(SNES snes, Vec X, Vec F, void *ctx)
47 {
48   DMSNES_Local *dmlocalsnes = (DMSNES_Local *)ctx;
49   DM            dm;
50   Vec           Xloc, Floc;
51   PetscBool     transform;
52 
53   PetscFunctionBegin;
54   PetscValidHeaderSpecific(snes, SNES_CLASSID, 1);
55   PetscValidHeaderSpecific(X, VEC_CLASSID, 2);
56   PetscValidHeaderSpecific(F, VEC_CLASSID, 3);
57   PetscCall(SNESGetDM(snes, &dm));
58   PetscCall(DMGetLocalVector(dm, &Xloc));
59   PetscCall(DMGetLocalVector(dm, &Floc));
60   PetscCall(VecZeroEntries(Xloc));
61   PetscCall(VecZeroEntries(Floc));
62   /* Non-conforming routines needs boundary values before G2L */
63   if (dmlocalsnes->boundarylocal) PetscCall((*dmlocalsnes->boundarylocal)(dm, Xloc, dmlocalsnes->boundarylocalctx));
64   PetscCall(DMGlobalToLocalBegin(dm, X, INSERT_VALUES, Xloc));
65   PetscCall(DMGlobalToLocalEnd(dm, X, INSERT_VALUES, Xloc));
66   /* Need to reset boundary values if we transformed */
67   PetscCall(DMHasBasisTransform(dm, &transform));
68   if (transform && dmlocalsnes->boundarylocal) PetscCall((*dmlocalsnes->boundarylocal)(dm, Xloc, dmlocalsnes->boundarylocalctx));
69   CHKMEMQ;
70   PetscCall((*dmlocalsnes->residuallocal)(dm, Xloc, Floc, dmlocalsnes->residuallocalctx));
71   CHKMEMQ;
72   PetscCall(VecZeroEntries(F));
73   PetscCall(DMLocalToGlobalBegin(dm, Floc, ADD_VALUES, F));
74   PetscCall(DMLocalToGlobalEnd(dm, Floc, ADD_VALUES, F));
75   PetscCall(DMRestoreLocalVector(dm, &Floc));
76   PetscCall(DMRestoreLocalVector(dm, &Xloc));
77   {
78     char        name[PETSC_MAX_PATH_LEN];
79     char        oldname[PETSC_MAX_PATH_LEN];
80     const char *tmp;
81     PetscInt    it;
82 
83     PetscCall(SNESGetIterationNumber(snes, &it));
84     PetscCall(PetscSNPrintf(name, PETSC_MAX_PATH_LEN, "Solution, Iterate %d", (int)it));
85     PetscCall(PetscObjectGetName((PetscObject)X, &tmp));
86     PetscCall(PetscStrncpy(oldname, tmp, PETSC_MAX_PATH_LEN - 1));
87     PetscCall(PetscObjectSetName((PetscObject)X, name));
88     PetscCall(VecViewFromOptions(X, (PetscObject)snes, "-dmsnes_solution_vec_view"));
89     PetscCall(PetscObjectSetName((PetscObject)X, oldname));
90     PetscCall(PetscSNPrintf(name, PETSC_MAX_PATH_LEN, "Residual, Iterate %d", (int)it));
91     PetscCall(PetscObjectSetName((PetscObject)F, name));
92     PetscCall(VecViewFromOptions(F, (PetscObject)snes, "-dmsnes_residual_vec_view"));
93   }
94   PetscFunctionReturn(PETSC_SUCCESS);
95 }
96 
97 static PetscErrorCode SNESComputeJacobian_DMLocal(SNES snes, Vec X, Mat A, Mat B, void *ctx)
98 {
99   DMSNES_Local *dmlocalsnes = (DMSNES_Local *)ctx;
100   DM            dm;
101   Vec           Xloc;
102   PetscBool     transform;
103 
104   PetscFunctionBegin;
105   PetscCall(SNESGetDM(snes, &dm));
106   if (dmlocalsnes->jacobianlocal) {
107     PetscCall(DMGetLocalVector(dm, &Xloc));
108     PetscCall(VecZeroEntries(Xloc));
109     /* Non-conforming routines needs boundary values before G2L */
110     if (dmlocalsnes->boundarylocal) PetscCall((*dmlocalsnes->boundarylocal)(dm, Xloc, dmlocalsnes->boundarylocalctx));
111     PetscCall(DMGlobalToLocalBegin(dm, X, INSERT_VALUES, Xloc));
112     PetscCall(DMGlobalToLocalEnd(dm, X, INSERT_VALUES, Xloc));
113     /* Need to reset boundary values if we transformed */
114     PetscCall(DMHasBasisTransform(dm, &transform));
115     if (transform && dmlocalsnes->boundarylocal) PetscCall((*dmlocalsnes->boundarylocal)(dm, Xloc, dmlocalsnes->boundarylocalctx));
116     CHKMEMQ;
117     PetscCall((*dmlocalsnes->jacobianlocal)(dm, Xloc, A, B, dmlocalsnes->jacobianlocalctx));
118     CHKMEMQ;
119     PetscCall(DMRestoreLocalVector(dm, &Xloc));
120   } else {
121     MatFDColoring fdcoloring;
122     PetscCall(PetscObjectQuery((PetscObject)dm, "DMDASNES_FDCOLORING", (PetscObject *)&fdcoloring));
123     if (!fdcoloring) {
124       ISColoring coloring;
125 
126       PetscCall(DMCreateColoring(dm, dm->coloringtype, &coloring));
127       PetscCall(MatFDColoringCreate(B, coloring, &fdcoloring));
128       PetscCall(ISColoringDestroy(&coloring));
129       switch (dm->coloringtype) {
130       case IS_COLORING_GLOBAL:
131         PetscCall(MatFDColoringSetFunction(fdcoloring, (PetscErrorCode(*)(void))SNESComputeFunction_DMLocal, dmlocalsnes));
132         break;
133       default:
134         SETERRQ(PetscObjectComm((PetscObject)snes), PETSC_ERR_SUP, "No support for coloring type '%s'", ISColoringTypes[dm->coloringtype]);
135       }
136       PetscCall(PetscObjectSetOptionsPrefix((PetscObject)fdcoloring, ((PetscObject)dm)->prefix));
137       PetscCall(MatFDColoringSetFromOptions(fdcoloring));
138       PetscCall(MatFDColoringSetUp(B, coloring, fdcoloring));
139       PetscCall(PetscObjectCompose((PetscObject)dm, "DMDASNES_FDCOLORING", (PetscObject)fdcoloring));
140       PetscCall(PetscObjectDereference((PetscObject)fdcoloring));
141 
142       /* The following breaks an ugly reference counting loop that deserves a paragraph. MatFDColoringApply() will call
143        * VecDuplicate() with the state Vec and store inside the MatFDColoring. This Vec will duplicate the Vec, but the
144        * MatFDColoring is composed with the DM. We dereference the DM here so that the reference count will eventually
145        * drop to 0. Note the code in DMDestroy() that exits early for a negative reference count. That code path will be
146        * taken when the PetscObjectList for the Vec inside MatFDColoring is destroyed.
147        */
148       PetscCall(PetscObjectDereference((PetscObject)dm));
149     }
150     PetscCall(MatFDColoringApply(B, fdcoloring, X, snes));
151   }
152   /* This will be redundant if the user called both, but it's too common to forget. */
153   if (A != B) {
154     PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
155     PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
156   }
157   PetscFunctionReturn(PETSC_SUCCESS);
158 }
159 
160 /*@C
161    DMSNESSetFunctionLocal - set a local residual evaluation function. This function is called with local vector
162       containing the local vector information PLUS ghost point information. It should compute a result for all local
163       elements and `DMSNES` will automatically accumulate the overlapping values.
164 
165    Logically Collective
166 
167    Input Parameters:
168 +  dm - `DM` to associate callback with
169 .  func - local residual evaluation
170 -  ctx - optional context for local residual evaluation
171 
172    Level: advanced
173 
174 .seealso: `DMSNESSetFunction()`, `DMDASNESSetJacobianLocal()`, `DMDACreate1d()`, `DMDACreate2d()`, `DMDACreate3d()`
175 @*/
176 PetscErrorCode DMSNESSetFunctionLocal(DM dm, PetscErrorCode (*func)(DM, Vec, Vec, void *), void *ctx)
177 {
178   DMSNES        sdm;
179   DMSNES_Local *dmlocalsnes;
180 
181   PetscFunctionBegin;
182   PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
183   PetscCall(DMGetDMSNESWrite(dm, &sdm));
184   PetscCall(DMLocalSNESGetContext(dm, sdm, &dmlocalsnes));
185 
186   dmlocalsnes->residuallocal    = func;
187   dmlocalsnes->residuallocalctx = ctx;
188 
189   PetscCall(DMSNESSetFunction(dm, SNESComputeFunction_DMLocal, dmlocalsnes));
190   if (!sdm->ops->computejacobian) { /* Call us for the Jacobian too, can be overridden by the user. */
191     PetscCall(DMSNESSetJacobian(dm, SNESComputeJacobian_DMLocal, dmlocalsnes));
192   }
193   PetscFunctionReturn(PETSC_SUCCESS);
194 }
195 
196 /*@C
197    DMSNESSetBoundaryLocal - set a local boundary value function. This function is called with local vector
198       containing the local vector information PLUS ghost point information. It should insert values into the local
199       vector that do not come from the global vector, such as essential boundary condition data.
200 
201    Logically Collective
202 
203    Input Parameters:
204 +  dm - `DM` to associate callback with
205 .  func - local boundary value evaluation
206 -  ctx - optional context for local boundary value evaluation
207 
208    Level: advanced
209 
210 .seealso: `DMSNESSetFunctionLocal()`, `DMDASNESSetJacobianLocal()`
211 @*/
212 PetscErrorCode DMSNESSetBoundaryLocal(DM dm, PetscErrorCode (*func)(DM, Vec, void *), void *ctx)
213 {
214   DMSNES        sdm;
215   DMSNES_Local *dmlocalsnes;
216 
217   PetscFunctionBegin;
218   PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
219   PetscCall(DMGetDMSNESWrite(dm, &sdm));
220   PetscCall(DMLocalSNESGetContext(dm, sdm, &dmlocalsnes));
221 
222   dmlocalsnes->boundarylocal    = func;
223   dmlocalsnes->boundarylocalctx = ctx;
224 
225   PetscFunctionReturn(PETSC_SUCCESS);
226 }
227 
228 /*@C
229    DMSNESSetJacobianLocal - set a local Jacobian evaluation function
230 
231    Logically Collective
232 
233    Input Parameters:
234 +  dm - DM to associate callback with
235 .  func - local Jacobian evaluation
236 -  ctx - optional context for local Jacobian evaluation
237 
238    Level: advanced
239 
240 .seealso: `DMSNESSetJacobian()`, `DMDASNESSetJacobian()`, `DMDACreate1d()`, `DMDACreate2d()`, `DMDACreate3d()`
241 @*/
242 PetscErrorCode DMSNESSetJacobianLocal(DM dm, PetscErrorCode (*func)(DM, Vec, Mat, Mat, void *), void *ctx)
243 {
244   DMSNES        sdm;
245   DMSNES_Local *dmlocalsnes;
246 
247   PetscFunctionBegin;
248   PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
249   PetscCall(DMGetDMSNESWrite(dm, &sdm));
250   PetscCall(DMLocalSNESGetContext(dm, sdm, &dmlocalsnes));
251 
252   dmlocalsnes->jacobianlocal    = func;
253   dmlocalsnes->jacobianlocalctx = ctx;
254 
255   PetscCall(DMSNESSetJacobian(dm, SNESComputeJacobian_DMLocal, dmlocalsnes));
256   PetscFunctionReturn(PETSC_SUCCESS);
257 }
258 
259 /*@C
260    DMSNESGetFunctionLocal - get the local residual evaluation function information set with `DMSNESSetFunctionLocal()`.
261 
262    Not Collective
263 
264    Input Parameter:
265 .  dm - `DM` with the associated callback
266 
267    Output Parameters:
268 +  func - local residual evaluation
269 -  ctx - context for local residual evaluation
270 
271    Level: beginner
272 
273 .seealso: `DMSNESSetFunction()`, `DMSNESSetFunctionLocal()`, `DMDASNESSetJacobianLocal()`, `DMDACreate1d()`, `DMDACreate2d()`, `DMDACreate3d()`
274 @*/
275 PetscErrorCode DMSNESGetFunctionLocal(DM dm, PetscErrorCode (**func)(DM, Vec, Vec, void *), void **ctx)
276 {
277   DMSNES        sdm;
278   DMSNES_Local *dmlocalsnes;
279 
280   PetscFunctionBegin;
281   PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
282   PetscCall(DMGetDMSNES(dm, &sdm));
283   PetscCall(DMLocalSNESGetContext(dm, sdm, &dmlocalsnes));
284   if (func) *func = dmlocalsnes->residuallocal;
285   if (ctx) *ctx = dmlocalsnes->residuallocalctx;
286   PetscFunctionReturn(PETSC_SUCCESS);
287 }
288 
289 /*@C
290    DMSNESGetBoundaryLocal - get the local boundary value function set with `DMSNESSetBoundaryLocal()`.
291 
292    Not Collective
293 
294    Input Parameter:
295 .  dm - `DM` with the associated callback
296 
297    Output Parameters:
298 +  func - local boundary value evaluation
299 -  ctx - context for local boundary value evaluation
300 
301    Level: intermediate
302 
303 .seealso: `DMSNESSetFunctionLocal()`, `DMSNESSetBoundaryLocal()`, `DMDASNESSetJacobianLocal()`
304 @*/
305 PetscErrorCode DMSNESGetBoundaryLocal(DM dm, PetscErrorCode (**func)(DM, Vec, void *), void **ctx)
306 {
307   DMSNES        sdm;
308   DMSNES_Local *dmlocalsnes;
309 
310   PetscFunctionBegin;
311   PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
312   PetscCall(DMGetDMSNES(dm, &sdm));
313   PetscCall(DMLocalSNESGetContext(dm, sdm, &dmlocalsnes));
314   if (func) *func = dmlocalsnes->boundarylocal;
315   if (ctx) *ctx = dmlocalsnes->boundarylocalctx;
316   PetscFunctionReturn(PETSC_SUCCESS);
317 }
318 
319 /*@C
320    DMSNESGetJacobianLocal - the local Jacobian evaluation function set with `DMSNESSetJacobianLocal()`.
321 
322    Logically Collective
323 
324    Input Parameter:
325 .  dm - `DM` with the associated callback
326 
327    Output Parameters:
328 +  func - local Jacobian evaluation
329 -  ctx - context for local Jacobian evaluation
330 
331    Level: beginner
332 
333 .seealso: `DMSNESSetJacobianLocal()`, `DMDASNESSetJacobian()`, `DMDACreate1d()`, `DMDACreate2d()`, `DMDACreate3d()`
334 @*/
335 PetscErrorCode DMSNESGetJacobianLocal(DM dm, PetscErrorCode (**func)(DM, Vec, Mat, Mat, void *), void **ctx)
336 {
337   DMSNES        sdm;
338   DMSNES_Local *dmlocalsnes;
339 
340   PetscFunctionBegin;
341   PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
342   PetscCall(DMGetDMSNES(dm, &sdm));
343   PetscCall(DMLocalSNESGetContext(dm, sdm, &dmlocalsnes));
344   if (func) *func = dmlocalsnes->jacobianlocal;
345   if (ctx) *ctx = dmlocalsnes->jacobianlocalctx;
346   PetscFunctionReturn(PETSC_SUCCESS);
347 }
348